pkg/aflow: add package

Add infrastructure for defining LLM agent workflows.
The workflows can be executed using the tools/syz-agent tool
from the command line, or programmatically.

Author: dvyukov

go.mod

@@ -18,6 +18,7 @@ require (
 	github.com/google/flatbuffers v25.9.23+incompatible
 	github.com/google/generative-ai-go v0.20.1
 	github.com/google/go-cmp v0.7.0
+	github.com/google/jsonschema-go v0.3.0
 	github.com/google/uuid v1.6.0
 	github.com/gorilla/handlers v1.5.2
 	github.com/ianlancetaylor/demangle v0.0.0-20250628045327-2d64ad6b7ec5
@@ -32,8 +33,10 @@ require (
 	golang.org/x/sync v0.17.0
 	golang.org/x/sys v0.37.0
 	golang.org/x/tools v0.38.0
+	google.golang.org/adk v0.0.0-20251024065725-37e04037159e
 	google.golang.org/api v0.252.0
 	google.golang.org/appengine/v2 v2.0.6
+	google.golang.org/genai v1.32.0
 	google.golang.org/genproto v0.0.0-20250603155806-513f23925822
 	google.golang.org/grpc v1.76.0
 	google.golang.org/protobuf v1.36.10
@@ -172,6 +175,7 @@ require (
 	github.com/googleapis/enterprise-certificate-proxy v0.3.6 // indirect
 	github.com/googleapis/gax-go/v2 v2.15.0 // indirect
 	github.com/gordonklaus/ineffassign v0.1.0 // indirect
+	github.com/gorilla/websocket v1.5.4-0.20250319132907-e064f32e3674 // indirect
 	github.com/gostaticanalysis/analysisutil v0.7.1 // indirect
 	github.com/gostaticanalysis/comment v1.5.0 // indirect
 	github.com/gostaticanalysis/forcetypeassert v0.2.0 // indirect
@@ -268,7 +272,7 @@ require (
 	github.com/spf13/afero v1.14.0 // indirect
 	github.com/spf13/cast v1.9.2 // indirect
 	github.com/spf13/cobra v1.9.1 // indirect
-	github.com/spf13/pflag v1.0.7 // indirect
+	github.com/spf13/pflag v1.0.10 // indirect
 	github.com/spf13/viper v1.20.1 // indirect
 	github.com/spiffe/go-spiffe/v2 v2.6.0 // indirect
 	github.com/ssgreg/nlreturn/v2 v2.2.1 // indirect
@@ -336,6 +340,8 @@ require (
 	k8s.io/utils v0.0.0-20251002143259-bc988d571ff4 // indirect
 	mvdan.cc/gofumpt v0.8.0 // indirect
 	mvdan.cc/unparam v0.0.0-20250301125049-0df0534333a4 // indirect
+	rsc.io/omap v1.2.0 // indirect
+	rsc.io/ordered v1.1.1 // indirect
 	sigs.k8s.io/json v0.0.0-20250730193827-2d320260d730 // indirect
 	sigs.k8s.io/randfill v1.0.0 // indirect
 	sigs.k8s.io/structured-merge-diff/v6 v6.3.0 // indirect

go.sum

@@ -1033,6 +1033,8 @@ github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeN
 github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
 github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
 github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
+github.com/google/jsonschema-go v0.3.0 h1:6AH2TxVNtk3IlvkkhjrtbUc4S8AvO0Xii0DxIygDg+Q=
+github.com/google/jsonschema-go v0.3.0/go.mod h1:r5quNTdLOYEz95Ru18zA0ydNbBuYoo9tgaYcxEYhJVE=
 github.com/google/keep-sorted v0.6.1 h1:LNEdDKYxoXOrn4ZXC+FdUfJCVbUjhb2QPIBs5XISXCI=
 github.com/google/keep-sorted v0.6.1/go.mod h1:JYy9vljs7P8b3QdPOQkywA+4u36FUHwsNITZIpJyPkE=
 github.com/google/martian v2.1.0+incompatible h1:/CP5g8u/VJHijgedC/Legn3BAbAaWPgecwXBIDzw5no=
@@ -1095,6 +1097,8 @@ github.com/gordonklaus/ineffassign v0.1.0 h1:y2Gd/9I7MdY1oEIt+n+rowjBNDcLQq3RsH5
 github.com/gordonklaus/ineffassign v0.1.0/go.mod h1:Qcp2HIAYhR7mNUVSIxZww3Guk4it82ghYcEXIAk+QT0=
 github.com/gorilla/handlers v1.5.2 h1:cLTUSsNkgcwhgRqvCNmdbRWG0A3N4F+M2nWKdScwyEE=
 github.com/gorilla/handlers v1.5.2/go.mod h1:dX+xVpaxdSw+q0Qek8SSsl3dfMk3jNddUkMzo0GtH0w=
+github.com/gorilla/websocket v1.5.4-0.20250319132907-e064f32e3674 h1:JeSE6pjso5THxAzdVpqr6/geYxZytqFMBCOtn/ujyeo=
+github.com/gorilla/websocket v1.5.4-0.20250319132907-e064f32e3674/go.mod h1:r4w70xmWCQKmi1ONH4KIaBptdivuRPyosB9RmPlGEwA=
 github.com/gostaticanalysis/analysisutil v0.7.1 h1:ZMCjoue3DtDWQ5WyU16YbjbQEQ3VuzwxALrpYd+HeKk=
 github.com/gostaticanalysis/analysisutil v0.7.1/go.mod h1:v21E3hY37WKMGSnbsw2S/ojApNWb6C1//mXO48CXbVc=
 github.com/gostaticanalysis/comment v1.4.1/go.mod h1:ih6ZxzTHLdadaiSnF5WY3dxUoXfXAlTaRzuaNDlSado=
@@ -1396,8 +1400,8 @@ github.com/spf13/cobra v1.9.1 h1:CXSaggrXdbHK9CF+8ywj8Amf7PBRmPCOJugH954Nnlo=
 github.com/spf13/cobra v1.9.1/go.mod h1:nDyEzZ8ogv936Cinf6g1RU9MRY64Ir93oCnqb9wxYW0=
 github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
 github.com/spf13/pflag v1.0.6/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
-github.com/spf13/pflag v1.0.7 h1:vN6T9TfwStFPFM5XzjsvmzZkLuaLX+HS+0SeFLRgU6M=
-github.com/spf13/pflag v1.0.7/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
+github.com/spf13/pflag v1.0.10 h1:4EBh2KAYBwaONj6b2Ye1GiHfwjqyROoF4RwYO+vPwFk=
+github.com/spf13/pflag v1.0.10/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
 github.com/spf13/viper v1.20.1 h1:ZMi+z/lvLyPSCoNtFCpqjy0S4kPbirhpTMwl8BkW9X4=
 github.com/spf13/viper v1.20.1/go.mod h1:P9Mdzt1zoHIG8m2eZQinpiBjo6kCmZSKBClNNqjJvu4=
 github.com/spiffe/go-spiffe/v2 v2.6.0 h1:l+DolpxNWYgruGQVV0xsfeya3CsC7m8iBzDnMpsbLuo=
@@ -1966,6 +1970,8 @@ gonum.org/v1/netlib v0.0.0-20190313105609-8cb42192e0e0/go.mod h1:wa6Ws7BG/ESfp6d
 gonum.org/v1/plot v0.0.0-20190515093506-e2840ee46a6b/go.mod h1:Wt8AAjI+ypCyYX3nZBvf6cAIx93T+c/OS2HFAYskSZc=
 gonum.org/v1/plot v0.9.0/go.mod h1:3Pcqqmp6RHvJI72kgb8fThyUnav364FOsdDo2aGW5lY=
 gonum.org/v1/plot v0.10.1/go.mod h1:VZW5OlhkL1mysU9vaqNHnsy86inf6Ot+jB3r+BczCEo=
+google.golang.org/adk v0.0.0-20251024065725-37e04037159e h1:5ZdMBo/z3nbrSjB/s5ajewcLDKIuBDEUOmiSimNgxk0=
+google.golang.org/adk v0.0.0-20251024065725-37e04037159e/go.mod h1:ToWN2tT+8VZb7zxH4trB9Vxetlvig3u5uwtQpNCM8ZM=
 google.golang.org/api v0.4.0/go.mod h1:8k5glujaEP+g9n7WNsDg8QP6cUVNI86fCNMcbazEtwE=
 google.golang.org/api v0.7.0/go.mod h1:WtwebWUNSVBH/HAw79HIFXZNqEvBhG+Ra+ax0hx3E3M=
 google.golang.org/api v0.8.0/go.mod h1:o4eAsZoiT+ibD93RtjEohWalFOjRDx6CVaqeizhEnKg=
@@ -2034,6 +2040,8 @@ google.golang.org/appengine v1.6.6/go.mod h1:8WjMMxjGQR8xUklV/ARdw2HLXBOI7O7uCID
 google.golang.org/appengine v1.6.7/go.mod h1:8WjMMxjGQR8xUklV/ARdw2HLXBOI7O7uCIDZVag1xfc=
 google.golang.org/appengine/v2 v2.0.6 h1:LvPZLGuchSBslPBp+LAhihBeGSiRh1myRoYK4NtuBIw=
 google.golang.org/appengine/v2 v2.0.6/go.mod h1:WoEXGoXNfa0mLvaH5sV3ZSGXwVmy8yf7Z1JKf3J3wLI=
+google.golang.org/genai v1.32.0 h1:kku/m3kWOncjnw8EIa2sgmrPLhaxFHaP+uqOq5ZckvI=
+google.golang.org/genai v1.32.0/go.mod h1:7pAilaICJlQBonjKKJNhftDFv3SREhZcTe9F6nRcjbg=
 google.golang.org/genproto v0.0.0-20180817151627-c66870c02cf8/go.mod h1:JiN7NxoALGmiZfu7CAH4rXhgtRTLTxftemlI0sWmxmc=
 google.golang.org/genproto v0.0.0-20190307195333-5fe7a883aa19/go.mod h1:VzzqZJRnGkLBvHegQrXjBqPurQTc5/KpmUdxsrq26oE=
 google.golang.org/genproto v0.0.0-20190418145605-e7d98fc518a7/go.mod h1:VzzqZJRnGkLBvHegQrXjBqPurQTc5/KpmUdxsrq26oE=
@@ -2312,6 +2320,10 @@ mvdan.cc/gofumpt v0.8.0/go.mod h1:vEYnSzyGPmjvFkqJWtXkh79UwPWP9/HMxQdGEXZHjpg=
 mvdan.cc/unparam v0.0.0-20250301125049-0df0534333a4 h1:WjUu4yQoT5BHT1w8Zu56SP8367OuBV5jvo+4Ulppyf8=
 mvdan.cc/unparam v0.0.0-20250301125049-0df0534333a4/go.mod h1:rthT7OuvRbaGcd5ginj6dA2oLE7YNlta9qhBNNdCaLE=
 rsc.io/binaryregexp v0.2.0/go.mod h1:qTv7/COck+e2FymRvadv62gMdZztPaShugOCi3I+8D8=
+rsc.io/omap v1.2.0 h1:c1M8jchnHbzmJALzGLclfH3xDWXrPxSUHXzH5C+8Kdw=
+rsc.io/omap v1.2.0/go.mod h1:C8pkI0AWexHopQtZX+qiUeJGzvc8HkdgnsWK4/mAa00=
+rsc.io/ordered v1.1.1 h1:1kZM6RkTmceJgsFH/8DLQvkCVEYomVDJfBRLT595Uak=
+rsc.io/ordered v1.1.1/go.mod h1:evAi8739bWVBRG9aaufsjVc202+6okf8u2QeVL84BCM=
 rsc.io/pdf v0.1.1/go.mod h1:n8OzWcQ6Sp37PL01nO98y4iUCRdTGarVfzxY20ICaU4=
 rsc.io/quote/v3 v3.1.0/go.mod h1:yEA65RcK8LyAZtP9Kv3t0HmxON59tX3rD+tICJqUlj0=
 rsc.io/sampler v1.3.0/go.mod h1:T1hPZKmBbMNahiBKFy5HrXp6adAjACjK9JXDnKaTXpA=

pkg/aflow/action.go

@@ -0,0 +1,88 @@
+// Copyright 2025 syzkaller project authors. All rights reserved.
+// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.
+
+package aflow
+
+import (
+	"fmt"
+	"maps"
+	"slices"
+
+	"google.golang.org/adk/agent"
+	"google.golang.org/adk/agent/workflowagents/sequentialagent"
+	"google.golang.org/adk/model"
+	"google.golang.org/genai"
+)
+
+type Action interface {
+	create(*createContext) (agent.Agent, error)
+	verify(*verifyContext)
+}
+
+type Pipeline struct {
+	// For logging/debugging.
+	Name string
+	// These actions are invoked sequentially,
+	// but dataflow across actions is specified by their use
+	// of variables in args/instructions/prompts.
+	Actions []Action
+}
+
+type createContext struct {
+	llm           model.LLM
+	contentConfig *genai.GenerateContentConfig
+}
+
+type verifyContext struct {
+	state map[string]bool
+	err   error
+}
+
+func (vctx *verifyContext) errorf(who, msg string, args ...any) {
+	if vctx.err == nil {
+		vctx.err = fmt.Errorf(fmt.Sprintf("action %v: %v", who, msg), args...)
+	}
+}
+
+func (vctx *verifyContext) requireNotEmpty(who, name, value string) {
+	if value == "" {
+		vctx.errorf(who, "%v must not be empty", name)
+	}
+}
+
+func (vctx *verifyContext) requireInput(who, name string) {
+	if !vctx.state[name] {
+		vctx.errorf(who, "no input %v, available inputs: %v",
+			name, slices.Collect(maps.Keys(vctx.state)))
+	}
+}
+
+func (vctx *verifyContext) provideOutput(who, name string, unique bool) {
+	if unique && vctx.state[name] {
+		vctx.errorf(who, "output %v is already set", name)
+	}
+	vctx.state[name] = true
+}
+
+func (p *Pipeline) create(cctx *createContext) (agent.Agent, error) {
+	var agents []agent.Agent
+	for _, sub := range p.Actions {
+		subAgent, err := sub.create(cctx)
+		if err != nil {
+			return nil, err
+		}
+		agents = append(agents, subAgent)
+	}
+	return sequentialagent.New(sequentialagent.Config{
+		AgentConfig: agent.Config{
+			Name:      p.Name,
+			SubAgents: agents,
+		},
+	})
+}
+
+func (p *Pipeline) verify(vctx *verifyContext) {
+	for _, a := range p.Actions {
+		a.verify(vctx)
+	}
+}

pkg/aflow/execute.go

@@ -0,0 +1,124 @@
+// Copyright 2025 syzkaller project authors. All rights reserved.
+// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.
+
+package aflow
+
+import (
+	"context"
+
+	"google.golang.org/adk/agent"
+	"google.golang.org/adk/model"
+	"google.golang.org/adk/model/gemini"
+	"google.golang.org/adk/runner"
+	"google.golang.org/adk/session"
+	"google.golang.org/genai"
+)
+
+const (
+	largeModelName = "gemini-2.5-pro"
+	smallModelName = "gemini-2.0-flash-lite" // cheaper/faster (for smoke testing)
+	maxLLMCalls    = 100
+)
+
+type Callback interface {
+	OnRequest(agentName string, req *model.LLMRequest) error
+	OnEvent(*session.Event) error
+}
+
+type Context interface {
+	Workdir() string
+	onRequest(agentName string, req *model.LLMRequest) error
+}
+
+type executeContext struct {
+	callback Callback
+	workdir  string
+}
+
+func (ectx *executeContext) Workdir() string {
+	return ectx.workdir
+}
+
+func (ectx *executeContext) onRequest(agentName string, req *model.LLMRequest) error {
+	return ectx.callback.OnRequest(agentName, req)
+}
+
+type executeContextKeyType int
+
+var executeContextKey executeContextKeyType = 0
+
+func (flow *Flow) Execute(ctx context.Context, largeModel bool, workdir string,
+	inputs any, events []*session.Event, cb Callback) (any, error) {
+	modelName := smallModelName
+	var thinkingConfig *genai.ThinkingConfig
+	if largeModel {
+		modelName = largeModelName
+		thinkingConfig = &genai.ThinkingConfig{
+			IncludeThoughts: true,
+			ThinkingBudget:  genai.Ptr[int32](-1),
+		}
+	}
+	ectx := &executeContext{
+		callback: cb,
+		workdir:  workdir,
+	}
+	ctx = context.WithValue(ctx, executeContextKey, ectx)
+	llm, err := gemini.NewModel(ctx, modelName, &genai.ClientConfig{})
+	if err != nil {
+		return nil, err
+	}
+	cctx := &createContext{
+		llm: llm,
+		contentConfig: &genai.GenerateContentConfig{
+			Temperature:    genai.Ptr[float32](0),
+			ThinkingConfig: thinkingConfig,
+		},
+	}
+	root, err := flow.Root.create(cctx)
+	if err != nil {
+		return nil, err
+	}
+	sessions := session.InMemoryService()
+	const (
+		userID    = "user"
+		sessionID = "session"
+	)
+	createReq := &session.CreateRequest{
+		AppName:   flow.Name,
+		UserID:    userID,
+		SessionID: sessionID,
+	}
+	createResp, err := sessions.Create(ctx, createReq)
+	if err != nil {
+		return nil, err
+	}
+	session := createResp.Session
+	if flow.convertInputs(session.State(), inputs); err != nil {
+		return nil, err
+	}
+	for _, ev := range events {
+		if err := sessions.AppendEvent(ctx, session, ev); err != nil {
+			return nil, err
+		}
+	}
+	r, err := runner.New(runner.Config{
+		AppName:        flow.Name,
+		Agent:          root,
+		SessionService: sessions,
+	})
+	if err != nil {
+		return nil, err
+	}
+	cfg := agent.RunConfig{
+		MaxLLMCalls: maxLLMCalls,
+	}
+	for ev, err := range r.Run(ctx, userID, sessionID, nil, cfg) {
+		if err != nil {
+			return nil, err
+		}
+		if err := cb.OnEvent(ev); err != nil {
+			return nil, err
+		}
+	}
+	return flow.extractOutputs(session.State())
+}