chore: add interface assertion for readability

agent/agent.go

@@ -145,6 +145,8 @@ type agent struct {
 	afterAgentCallbacks  []AfterAgentCallback
 }
 
+var _ Agent = (*agent)(nil)
+
 func (a *agent) Name() string {
 	return a.name
 }
@@ -440,6 +442,8 @@ type invocationContext struct {
 	endInvocation bool
 }
 
+var _ InvocationContext = (*invocationContext)(nil)
+
 func (c *invocationContext) Agent() Agent {
 	return c.agent
 }
@@ -499,5 +503,3 @@ type pluginManager interface {
 	RunBeforeAgentCallback(cctx CallbackContext) (*genai.Content, error)
 	RunAfterAgentCallback(cctx CallbackContext) (*genai.Content, error)
 }
-
-var _ InvocationContext = (*invocationContext)(nil)

agent/agent_test.go

@@ -277,4 +277,6 @@ type mockSession struct {
 	sessionID string
 }
 
+var _ session.Session = (*mockSession)(nil)
+
 func (m *mockSession) ID() string { return m.sessionID }

agent/llmagent/llmagent.go

@@ -342,6 +342,8 @@ type llmAgent struct {
 	outputSchema *genai.Schema
 }
 
+var _ agentinternal.Agent = (*llmAgent)(nil)
+
 type agentState = agentinternal.State
 
 func (a *llmAgent) run(ctx agent.InvocationContext) iter.Seq2[*session.Event, error] {

agent/llmagent/state_agent_test.go

@@ -41,6 +41,8 @@ type FakeLLM struct {
 	GenerateContentFunc func(ctx context.Context, req *model.LLMRequest, stream bool) (model.LLMResponse, error)
 }
 
+var _ model.LLM = (*FakeLLM)(nil)
+
 func (f *FakeLLM) Name() string {
 	return "fake-llm"
 }

agent/loader.go

@@ -28,17 +28,13 @@ type Loader interface {
 	RootAgent() Agent
 }
 
-// multiLoader should be used when you have multiple agents
-type multiLoader struct {
-	agentMap map[string]Agent
-	root     Agent
-}
-
 // singleLoader should be used when you have only one agent
 type singleLoader struct {
 	root Agent
 }
 
+var _ Loader = (*singleLoader)(nil)
+
 // NewSingleLoader returns a loader with only one agent, which becomes the root agent
 func NewSingleLoader(a Agent) Loader {
 	return &singleLoader{root: a}
@@ -65,6 +61,14 @@ func (s *singleLoader) RootAgent() Agent {
 	return s.root
 }
 
+// multiLoader should be used when you have multiple agents
+type multiLoader struct {
+	agentMap map[string]Agent
+	root     Agent
+}
+
+var _ Loader = (*multiLoader)(nil)
+
 // NewMultiLoader returns a new AgentLoader with the given root Agent and other agents.
 // Returns an error if more than one agent (including root) shares the same name
 func NewMultiLoader(root Agent, agents ...Agent) (Loader, error) {

agent/loader_test.go

@@ -21,12 +21,12 @@ import (
 	"google.golang.org/adk/session"
 )
 
-var _ Agent = (*testAgent)(nil)
-
 type testAgent struct {
 	name string
 }
 
+var _ Agent = (*testAgent)(nil)
+
 func (a *testAgent) Name() string {
 	return a.name
 }

agent/remoteagent/a2a_e2e_test.go

@@ -737,6 +737,8 @@ type llmStub struct {
 	generateContent func(ctx context.Context, req *model.LLMRequest, stream bool) iter.Seq2[*model.LLMResponse, error]
 }
 
+var _ model.LLM = (*llmStub)(nil)
+
 func (d *llmStub) Name() string {
 	return d.name
 }

agent/workflowagents/loopagent/agent_test.go

@@ -334,6 +334,8 @@ type FakeLLM struct {
 	skipSummarization bool
 }
 
+var _ model.LLM = (*FakeLLM)(nil)
+
 func (f *FakeLLM) Name() string {
 	return "fake-llm"
 }

agent/workflowagents/sequentialagent/agent_test.go

@@ -299,6 +299,8 @@ type FakeLLM struct {
 	callCounter int
 }
 
+var _ model.LLM = (*FakeLLM)(nil)
+
 func (f *FakeLLM) Name() string {
 	return "fake-llm"
 }

artifact/gcsartifact/gcs_client.go

@@ -59,6 +59,8 @@ type gcsClientWrapper struct {
 	client *storage.Client
 }
 
+var _ gcsClient = (*gcsClientWrapper)(nil)
+
 // Bucket returns a gcsBucketWrapper that satisfies the gcsBucket interface.
 func (w *gcsClientWrapper) bucket(name string) gcsBucket {
 	return &gcsBucketWrapper{
@@ -71,6 +73,8 @@ type gcsBucketWrapper struct {
 	bucket *storage.BucketHandle
 }
 
+var _ gcsBucket = (*gcsBucketWrapper)(nil)
+
 // Object returns a gcsObjectWrapper that satisfies the gcsObject interface.
 func (w *gcsBucketWrapper) object(name string) gcsObject {
 	objectHandle := w.bucket.Object(name)
@@ -92,6 +96,8 @@ type gcsObjectWrapper struct {
 	object *storage.ObjectHandle
 }
 
+var _ gcsObject = (*gcsObjectWrapper)(nil)
+
 // NewWriter implements the gcsObject interface for gcsObjectWrapper.
 func (w *gcsObjectWrapper) newWriter(ctx context.Context) gcsWriter {
 	return &gcsWriterWrapper{w: w.object.NewWriter(ctx)}
@@ -117,6 +123,8 @@ type gcsObjectIteratorWrapper struct {
 	iter *storage.ObjectIterator
 }
 
+var _ gcsObjectIterator = (*gcsObjectIteratorWrapper)(nil)
+
 func (w *gcsObjectIteratorWrapper) next() (*storage.ObjectAttrs, error) {
 	return w.iter.Next()
 }
@@ -126,6 +134,8 @@ type gcsWriterWrapper struct {
 	w *storage.Writer
 }
 
+var _ gcsWriter = (*gcsWriterWrapper)(nil)
+
 func (g *gcsWriterWrapper) Write(p []byte) (n int, err error) {
 	return g.w.Write(p)
 }
@@ -137,11 +147,3 @@ func (g *gcsWriterWrapper) Close() error {
 func (g *gcsWriterWrapper) SetContentType(cType string) {
 	g.w.ContentType = cType
 }
-
-var (
-	_ gcsClient         = (*gcsClientWrapper)(nil)
-	_ gcsBucket         = (*gcsBucketWrapper)(nil)
-	_ gcsObject         = (*gcsObjectWrapper)(nil)
-	_ gcsObjectIterator = (*gcsObjectIteratorWrapper)(nil)
-	_ gcsWriter         = (*gcsWriterWrapper)(nil)
-)

artifact/gcsartifact/gcs_test.go

@@ -55,6 +55,8 @@ type fakeClient struct {
 	inMemoryBucket gcsBucket
 }
 
+var _ gcsClient = (*fakeClient)(nil)
+
 func newFakeClient() gcsClient {
 	return &fakeClient{
 		inMemoryBucket: &fakeBucket{
@@ -74,6 +76,8 @@ type fakeBucket struct {
 	objectsMap map[string]*fakeObject
 }
 
+var _ gcsBucket = (*fakeBucket)(nil)
+
 // Object returns a fake object from the in-memory store.
 func (f *fakeBucket) object(name string) gcsObject {
 	f.mu.Lock()
@@ -116,6 +120,8 @@ type fakeObject struct {
 	contentType string
 }
 
+var _ gcsObject = (*fakeObject)(nil)
+
 // NewWriter returns a fake writer that stores data in memory.
 func (f *fakeObject) newWriter(ctx context.Context) gcsWriter {
 	f.mu.Lock()
@@ -160,6 +166,8 @@ type fakeWriter struct {
 	contentType string
 }
 
+var _ gcsWriter = (*fakeWriter)(nil)
+
 func (w *fakeWriter) Write(p []byte) (n int, err error) {
 	return w.buffer.Write(p)
 }
@@ -184,6 +192,8 @@ type fakeObjectIterator struct {
 	index   int
 }
 
+var _ gcsObjectIterator = (*fakeObjectIterator)(nil)
+
 // Next implements the iterator pattern.
 // It returns the next object in the slice or an iterator.Done error.
 func (i *fakeObjectIterator) next() (*storage.ObjectAttrs, error) {
@@ -194,11 +204,3 @@ func (i *fakeObjectIterator) next() (*storage.ObjectAttrs, error) {
 	i.index++
 	return &storage.ObjectAttrs{Name: obj.name, ContentType: obj.contentType}, nil
 }
-
-var (
-	_ gcsClient         = (*fakeClient)(nil)
-	_ gcsBucket         = (*fakeBucket)(nil)
-	_ gcsObject         = (*fakeObject)(nil)
-	_ gcsObjectIterator = (*fakeObjectIterator)(nil)
-	_ gcsWriter         = (*fakeWriter)(nil)
-)

artifact/gcsartifact/service.go

@@ -39,13 +39,15 @@ import (
 	"google.golang.org/adk/artifact"
 )
 
-// gcsService is a google cloud storage implementation of the Service.
+// gcsService is a google cloud storage implementation of the [artifact.Service].
 type gcsService struct {
 	bucketName    string
 	storageClient gcsClient
 	bucket        gcsBucket
 }
 
+var _ artifact.Service = (*gcsService)(nil)
+
 // NewService creates a Google Cloud Storage service for the specified bucket.
 func NewService(ctx context.Context, bucketName string, opts ...option.ClientOption) (artifact.Service, error) {
 	storageClient, err := storage.NewClient(ctx, opts...)

artifact/inmemory.go

@@ -39,6 +39,8 @@ type inMemoryService struct {
 	artifacts omap.Map[string, *genai.Part]
 }
 
+var _ Service = (*inMemoryService)(nil)
+
 // InMemoryService returns a new in-memory artifact service.
 func InMemoryService() Service {
 	return &inMemoryService{}
@@ -284,5 +286,3 @@ func (s *inMemoryService) Versions(ctx context.Context, req *VersionsRequest) (*
 	}
 	return &VersionsResponse{Versions: versions}, nil
 }
-
-var _ Service = (*inMemoryService)(nil)

internal/agent/state.go

@@ -14,7 +14,7 @@
 
 package agent
 
-// holds Agent internal state
+// Agent holds agent internal state.
 type Agent interface {
 	internal() *State
 }
@@ -24,6 +24,12 @@ type State struct {
 	Config    any
 }
 
+var _ Agent = (*State)(nil)
+
+func (s *State) internal() *State { return s }
+
+func Reveal(a Agent) *State { return a.internal() }
+
 type Type string
 
 const (
@@ -33,7 +39,3 @@ const (
 	TypeParallelAgent   Type = "ParallelAgent"
 	TypeCustomAgent     Type = "CustomAgent"
 )
-
-func (s *State) internal() *State { return s }
-
-func Reveal(a Agent) *State { return a.internal() }

internal/context/invocation_context.go

@@ -54,6 +54,8 @@ type InvocationContext struct {
 	params InvocationContextParams
 }
 
+var _ agent.InvocationContext = (*InvocationContext)(nil)
+
 func (c *InvocationContext) Artifacts() agent.Artifacts {
 	return c.params.Artifacts
 }
@@ -99,5 +101,3 @@ func (c *InvocationContext) WithContext(ctx context.Context) agent.InvocationCon
 	newCtx.Context = ctx
 	return &newCtx
 }
-
-var _ agent.InvocationContext = (*InvocationContext)(nil)

internal/llminternal/agent.go

@@ -22,11 +22,13 @@ import (
 	"google.golang.org/adk/tool"
 )
 
-// holds LLMAgent internal state
+// Agent holds LLMAgent internal state.
 type Agent interface {
 	internal() *State
 }
 
+type InstructionProvider func(ctx agent.ReadonlyContext) (string, error)
+
 type State struct {
 	Model model.LLM
 
@@ -51,7 +53,7 @@ type State struct {
 	OutputKey string
 }
 
-type InstructionProvider func(ctx agent.ReadonlyContext) (string, error)
+var _ Agent = (*State)(nil)
 
 func (s *State) internal() *State { return s }
 

internal/llminternal/agent_transfer.go

@@ -97,22 +97,28 @@ func AgentTransferRequestProcessor(ctx agent.InvocationContext, req *model.LLMRe
 
 type TransferToAgentTool struct{}
 
-// Description implements tool.Tool.
+var (
+	_ toolinternal.FunctionTool     = (*TransferToAgentTool)(nil)
+	_ toolinternal.RequestProcessor = (*TransferToAgentTool)(nil)
+)
+
+// Description implements [toolinternal.FunctionTool].
 func (t *TransferToAgentTool) Description() string {
 	return `Transfer the question to another agent.
 This tool hands off control to another agent when it's more suitable to answer the user's question according to the agent's description.`
 }
 
-// Name implements tool.Tool.
+// Name implements [toolinternal.FunctionTool].
 func (t *TransferToAgentTool) Name() string {
 	return "transfer_to_agent"
 }
 
-// IsLongRunning implements tool.Tool.
+// IsLongRunning implements [toolinternal.FunctionTool].
 func (t *TransferToAgentTool) IsLongRunning() bool {
 	return false
 }
 
+// Declaration implements [toolinternal.FunctionTool].
 func (t *TransferToAgentTool) Declaration() *genai.FunctionDeclaration {
 	return &genai.FunctionDeclaration{
 		Name:        t.Name(),
@@ -130,12 +136,12 @@ func (t *TransferToAgentTool) Declaration() *genai.FunctionDeclaration {
 	}
 }
 
-// ProcessRequest implements types.Tool.
+// ProcessRequest implements [toolinternal.RequestProcessor].
 func (t *TransferToAgentTool) ProcessRequest(ctx tool.Context, req *model.LLMRequest) error {
 	return appendTools(req, t)
 }
 
-// Run implements types.Tool.
+// Run implements [toolinternal.FunctionTool].
 func (t *TransferToAgentTool) Run(ctx tool.Context, args any) (map[string]any, error) {
 	if args == nil {
 		return nil, fmt.Errorf("missing argument")
@@ -152,8 +158,6 @@ func (t *TransferToAgentTool) Run(ctx tool.Context, args any) (map[string]any, e
 	return map[string]any{}, nil
 }
 
-var _ tool.Tool = (*TransferToAgentTool)(nil)
-
 func transferTargets(agent, parent agent.Agent) []agent.Agent {
 	targets := slices.Clone(agent.SubAgents())