ronantakizawa/codereview-bench huggingface dataset:
Repository: google/adk-go
File: cmd/launcher/web/a2a/a2a.go (Go)

Code before:
	"golang.org/x/net/http2/h2c"
	"google.golang.org/adk/adka2a"
	"google.golang.org/adk/cmd/launcher/adk"
	"google.golang.org/adk/cmd/launcher/web"
	"google.golang.org/adk/internal/cli/util"
	"google.golang.org/adk/runner"
	"google.golang.org/grpc"
	"google.golang.org/grpc/reflection"
)

type a2aConfig struct {
	agentURL string
}

type A2ALauncher struct {
	flags  *flag.FlagSet
	config *a2aConfig
}

// NewLauncher creates new a2a launcher. It extends Web launcher
func NewLauncher() web.WebSublauncher {
	config := &a2aConfig{}

	fs := flag.NewFlagSet("a2a", flag.ContinueOnError)

	fs.StringVar(&config.agentURL, "a2a_agent_url", "localhost:8080", "A2A gRPC host URL as advertized in the public agent card. It is used by A2A clients as a connection endpoint.")

	return &A2ALauncher{
		config: config,
		flags:  fs,
	}
}

func (a *A2ALauncher) CommandLineSyntax() string {
	return util.FormatFlagUsage(a.flags)
}

func (a *A2ALauncher) Keyword() string {
	return "a2a"
}

func (a *A2ALauncher) Parse(args []string) ([]string, error) {
	err := a.flags.Parse(args)
	if err != nil || !a.flags.Parsed() {
		return nil, fmt.Errorf("failed to parse a2a flags: %v", err)
	}
	restArgs := a.flags.Args()
	return restArgs, nil
}

// WrapHandlers implements web.WebSublauncher. Returns http handler which basing on content-type

Review comment (suggestion):
```suggestion
	fs.StringVar(&config.agentURL, "a2a_agent_url", "localhost:8080", "A2A gRPC host URL as advertised in the public agent card. It is used by A2A clients as a connection endpoint.")
```

Code after:
	"golang.org/x/net/http2/h2c"
	"google.golang.org/adk/adka2a"
	"google.golang.org/adk/cmd/launcher/adk"
	"google.golang.org/adk/cmd/launcher/web"
	"google.golang.org/adk/internal/cli/util"
	"google.golang.org/adk/runner"
	"google.golang.org/grpc"
	"google.golang.org/grpc/reflection"
)

type a2aConfig struct {
	agentURL string
}

type a2aLauncher struct {
	flags  *flag.FlagSet
	config *a2aConfig
}

// NewLauncher creates new a2a launcher. It extends Web launcher
func NewLauncher() web.WebSublauncher {
	config := &a2aConfig{}

	fs := flag.NewFlagSet("a2a", flag.ContinueOnError)

	fs.StringVar(&config.agentURL, "a2a_agent_url", "localhost:8080", "A2A gRPC host URL as advertised in the public agent card. It is used by A2A clients as a connection endpoint.")

	return &a2aLauncher{
		config: config,
		flags:  fs,
	}
}

func (a *a2aLauncher) CommandLineSyntax() string {
	return util.FormatFlagUsage(a.flags)
}

func (a *a2aLauncher) Keyword() string {
	return "a2a"
}

func (a *a2aLauncher) Parse(args []string) ([]string, error) {
	err := a.flags.Parse(args)
	if err != nil || !a.flags.Parsed() {
		return nil, fmt.Errorf("failed to parse a2a flags: %v", err)
	}
	restArgs := a.flags.Args()
	return restArgs, nil
}

// WrapHandlers implements web.WebSublauncher. Returns http handler which basing on content-type
Repository: google/adk-go
File: artifact/inmemory.go (Go)

Code before:
	"strings"
	"sync"

	"google.golang.org/genai"
	"rsc.io/omap"
	"rsc.io/ordered"
)

// inMemoryService is an in-memory implementation of the Service.
// It is primarily for testing and demonstration purposes.
type inMemoryService struct {
	mu sync.RWMutex
	// ordered(appName, userID, sessionID) -> session
	artifacts omap.Map[string, *genai.Part]
}

func InMemoryService() Service {
	return &inMemoryService{}
}

// fileHasUserNamespace checks if a filename indicates a user scoped artifact.
func fileHasUserNamespace(filename string) bool {
	return strings.HasPrefix(filename, "user:")
}

// userPathPart defines the stirng for the part of the path used by user scope files
const userPathPart = "user"

type artifactKey struct {
	AppName   string
	UserID    string
	SessionID string
	FileName  string
	Version   int64
}

func (ak artifactKey) Encode() string {
	return string(ordered.Encode(ak.AppName, ak.UserID, ak.SessionID, ak.FileName, ordered.Rev(ak.Version)))
}

func (ak *artifactKey) Decode(key string) error {
	var v ordered.Reverse[int64]
	err := ordered.Decode([]byte(key), &ak.AppName, &ak.UserID, &ak.SessionID, &ak.FileName, &v)
	if err != nil {
		return err
	}
	ak.Version = v.Value()
	return nil
}

// scan returns an iterator over all key-value pairs

Review comment (suggestion):
```suggestion
// userPathPart defines the string for the part of the path used by user scope files
```

Code after:
	"strings"
	"sync"

	"google.golang.org/genai"
	"rsc.io/omap"
	"rsc.io/ordered"
)

// inMemoryService is an in-memory implementation of the Service.
// It is primarily for testing and demonstration purposes.
type inMemoryService struct {
	mu sync.RWMutex
	// ordered(appName, userID, sessionID) -> session
	artifacts omap.Map[string, *genai.Part]
}

func InMemoryService() Service {
	return &inMemoryService{}
}

// fileHasUserNamespace checks if a filename indicates a user scoped artifact.
func fileHasUserNamespace(filename string) bool {
	return strings.HasPrefix(filename, "user:")
}

// userScopedArtifactKey defines the string for the part of the path used by user scope files
// this key is used to store user artifacts within user namespace (available to all sessions for given app+user_id)
const userScopedArtifactKey = "user"

type artifactKey struct {
	AppName   string
	UserID    string
	SessionID string
	FileName  string
	Version   int64
}

func (ak artifactKey) Encode() string {
	return string(ordered.Encode(ak.AppName, ak.UserID, ak.SessionID, ak.FileName, ordered.Rev(ak.Version)))
}

func (ak *artifactKey) Decode(key string) error {
	var v ordered.Reverse[int64]
	err := ordered.Decode([]byte(key), &ak.AppName, &ak.UserID, &ak.SessionID, &ak.FileName, &v)
	if err != nil {
		return err
	}
	ak.Version = v.Value()
	return nil
}
Repository: google/adk-go
File: artifact/inmemory.go (Go)

Code before:
	"sync"

	"google.golang.org/genai"
	"rsc.io/omap"
	"rsc.io/ordered"
)

// inMemoryService is an in-memory implementation of the Service.
// It is primarily for testing and demonstration purposes.
type inMemoryService struct {
	mu sync.RWMutex
	// ordered(appName, userID, sessionID) -> session
	artifacts omap.Map[string, *genai.Part]
}

func InMemoryService() Service {
	return &inMemoryService{}
}

// fileHasUserNamespace checks if a filename indicates a user scoped artifact.
func fileHasUserNamespace(filename string) bool {
	return strings.HasPrefix(filename, "user:")
}

// userPathPart defines the stirng for the part of the path used by user scope files
const userPathPart = "user"

type artifactKey struct {
	AppName   string
	UserID    string
	SessionID string
	FileName  string
	Version   int64
}

func (ak artifactKey) Encode() string {
	return string(ordered.Encode(ak.AppName, ak.UserID, ak.SessionID, ak.FileName, ordered.Rev(ak.Version)))
}

func (ak *artifactKey) Decode(key string) error {
	var v ordered.Reverse[int64]
	err := ordered.Decode([]byte(key), &ak.AppName, &ak.UserID, &ak.SessionID, &ak.FileName, &v)
	if err != nil {
		return err
	}
	ak.Version = v.Value()
	return nil
}

// scan returns an iterator over all key-value pairs
// in the range begin ≤ key ≤ end.

Review comment (suggestion):
nit: 
```suggestion
// this key is used to store user artifacts within user namespace (available to all sessions for given app+user_id)
const userScopedArtifactKey = "user"
```

Or even make it a func something like `getUserScopedArtifactKey(...) artifactKey`

Code after:
	"sync"

	"google.golang.org/genai"
	"rsc.io/omap"
	"rsc.io/ordered"
)

// inMemoryService is an in-memory implementation of the Service.
// It is primarily for testing and demonstration purposes.
type inMemoryService struct {
	mu sync.RWMutex
	// ordered(appName, userID, sessionID) -> session
	artifacts omap.Map[string, *genai.Part]
}

func InMemoryService() Service {
	return &inMemoryService{}
}

// fileHasUserNamespace checks if a filename indicates a user scoped artifact.
func fileHasUserNamespace(filename string) bool {
	return strings.HasPrefix(filename, "user:")
}

// userScopedArtifactKey defines the string for the part of the path used by user scope files
// this key is used to store user artifacts within user namespace (available to all sessions for given app+user_id)
const userScopedArtifactKey = "user"

type artifactKey struct {
	AppName   string
	UserID    string
	SessionID string
	FileName  string
	Version   int64
}

func (ak artifactKey) Encode() string {
	return string(ordered.Encode(ak.AppName, ak.UserID, ak.SessionID, ak.FileName, ordered.Rev(ak.Version)))
}

func (ak *artifactKey) Decode(key string) error {
	var v ordered.Reverse[int64]
	err := ordered.Decode([]byte(key), &ak.AppName, &ak.UserID, &ak.SessionID, &ak.FileName, &v)
	if err != nil {
		return err
	}
	ak.Version = v.Value()
	return nil
}

// scan returns an iterator over all key-value pairs
Repository: google/adk-go
File: internal/context/invocation_context.go (Go)

Code before:
}

func (c *InvocationContext) InvocationID() string {
	return c.invocationID
}

func (c *InvocationContext) Memory() agent.Memory {
	return c.params.Memory
}

func (c *InvocationContext) Session() session.Session {
	return c.params.Session
}

func (c *InvocationContext) UserContent() *genai.Content {
	return c.params.UserContent
}

func (c *InvocationContext) RunConfig() *agent.RunConfig {
	return c.params.RunConfig
}

func (c *InvocationContext) EndInvocation() bool {
	return c.params.EndInvocation
}
func (c *InvocationContext) SetEndInvocation(endInvocation bool) {
	c.params.EndInvocation = endInvocation
}


Review comment (suggestion):
Maybe let's keep the `EndInvocation() bool` and `Ended()`?

Because it's not possible to un-end invocation `EndInvocation(false)` has no sense IMO.

Code after:
}

func (c *InvocationContext) InvocationID() string {
	return c.invocationID
}

func (c *InvocationContext) Memory() agent.Memory {
	return c.params.Memory
}

func (c *InvocationContext) Session() session.Session {
	return c.params.Session
}

func (c *InvocationContext) UserContent() *genai.Content {
	return c.params.UserContent
}

func (c *InvocationContext) RunConfig() *agent.RunConfig {
	return c.params.RunConfig
}

func (c *InvocationContext) EndInvocation() {
	c.params.EndInvocation = true
}
func (c *InvocationContext) Ended() bool {
	return c.params.EndInvocation
}
Repository: google/adk-go
File: server/adka2a/events.go (Go)

Code before:
			fnCall := converted[i]
			if fnCall.FunctionCall == nil {
				// TODO(yarolegovich): log a warning
				continue
			}
			ids = append(ids, fnCall.FunctionCall.ID)
		}
	}
	return ids
}

func toGenAIRole(role a2a.MessageRole) genai.Role {
	if role == a2a.MessageRoleUser {
		return genai.RoleUser
	} else {
		return genai.RoleModel
	}
}

func toEventActions(event a2a.Event) session.EventActions {
	meta := event.Meta()
	var result session.EventActions
	result.Escalate, _ = meta[metadataEscalateKey].(bool)
	if transferToAgent, ok := meta[metadataTransferToAgentKey].(string); ok {
		result.TransferToAgent = transferToAgent
	}
	return result
}


Review comment (suggestion):
optional: maybe shorten this code to be similar to the line above?
```suggestion
	if transferToAgent, _ = meta[metadataTransferToAgentKey].(string)
```

Code after:
			fnCall := converted[i]
			if fnCall.FunctionCall == nil {
				// TODO(yarolegovich): log a warning
				continue
			}
			ids = append(ids, fnCall.FunctionCall.ID)
		}
	}
	return ids
}

func toGenAIRole(role a2a.MessageRole) genai.Role {
	if role == a2a.MessageRoleUser {
		return genai.RoleUser
	} else {
		return genai.RoleModel
	}
}

func toEventActions(event a2a.Event) session.EventActions {
	meta := event.Meta()
	var result session.EventActions
	result.Escalate, _ = meta[metadataEscalateKey].(bool)
	result.TransferToAgent, _ = meta[metadataTransferToAgentKey].(string)
	return result
}
Repository: google/adk-go
File: tool/mcptoolset/client.go (Go)

Code before:
	defer c.mu.Unlock()

	if c.session != nil {
		return c.session, nil
	}

	session, err := c.client.Connect(ctx, c.transport, nil)
	if err != nil {
		return nil, fmt.Errorf("failed to init MCP session: %w", err)
	}

	c.session = session
	return c.session, nil
}

func (c *connectionRefresher) refreshConnection(ctx context.Context) (*mcp.ClientSession, bool, error) {
	c.mu.Lock()
	defer c.mu.Unlock()

	// First, try ping to confirm connection is dead
	if c.session != nil {
		if err := c.session.Ping(ctx, &mcp.PingParams{}); err == nil {
			// Connection is actually alive, don't refresh
			return c.session, false, nil
		}
		c.session.Close()
		c.session = nil
	}

	session, err := c.client.Connect(ctx, c.transport, nil)
	if err != nil {
		return nil, false, fmt.Errorf("failed to refresh MCP session: %w", err)
	}

	c.session = session
	return c.session, true, nil
}

var _ MCPClient = (*connectionRefresher)(nil)


Review comment (bug):
Let's fix the lint failures. Let's log the error here (in the app code) and t.Fail/Error in the test code?

Code after:
	}
	return result, false, err
}

// shouldRefreshConnection returns true if the error indicates we should
// attempt to refresh the MCP connection.
func shouldRefreshConnection(err error) bool {
	for _, target := range refreshableErrors {
		if errors.Is(err, target) {
			return true
		}
	}
	// TODO: replace with mcp.ErrSessionMissing when it's available https://github.com/modelcontextprotocol/go-sdk/issues/715
	if strings.Contains(err.Error(), "session not found") {
		return true
	}
	return false
}

func (c *connectionRefresher) getSession(ctx context.Context) (*mcp.ClientSession, error) {
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.session != nil {
		return c.session, nil
	}

	session, err := c.client.Connect(ctx, c.transport, nil)
	if err != nil {
		return nil, fmt.Errorf("failed to init MCP session: %w", err)
	}

	c.session = session
	return c.session, nil
}

func (c *connectionRefresher) refreshConnection(ctx context.Context) (*mcp.ClientSession, error) {
	c.mu.Lock()
	defer c.mu.Unlock()

	// Ping to verify the connection is actually dead before reconnecting.
	// This handles the case where another goroutine already reconnected.
	if c.session != nil {
		if err := c.session.Ping(ctx, &mcp.PingParams{}); err == nil {
			return c.session, nil
		}
		if err := c.session.Close(); err != nil {
			log.Printf("failed to close MCP session: %v", err)
		}
		c.session = nil
	}
Repository: google/adk-go
File: internal/toolinternal/context.go (Go)

Code before:
	agent.CallbackContext
	invocationContext agent.InvocationContext
	functionCallID    string
	eventActions      *session.EventActions
	artifacts         *internalArtifacts
}

func (c *toolContext) Artifacts() agent.Artifacts {
	return c.artifacts
}

func (c *toolContext) FunctionCallID() string {
	return c.functionCallID
}

func (c *toolContext) Actions() *session.EventActions {
	return c.eventActions
}

func (c *toolContext) AgentName() string {
	return c.invocationContext.Agent().Name()
}

func (c *toolContext) SearchMemory(ctx context.Context, query string) (*memory.SearchResponse, error) {
	if c.invocationContext.Memory() == nil {
		return nil, fmt.Errorf("memroy service nil error")
	}
	return c.invocationContext.Memory().Search(ctx, query)
}


Review comment (suggestion):
```suggestion
		return nil, fmt.Errorf("memory service is not set")
```

Code after:
	agent.CallbackContext
	invocationContext agent.InvocationContext
	functionCallID    string
	eventActions      *session.EventActions
	artifacts         *internalArtifacts
}

func (c *toolContext) Artifacts() agent.Artifacts {
	return c.artifacts
}

func (c *toolContext) FunctionCallID() string {
	return c.functionCallID
}

func (c *toolContext) Actions() *session.EventActions {
	return c.eventActions
}

func (c *toolContext) AgentName() string {
	return c.invocationContext.Agent().Name()
}

func (c *toolContext) SearchMemory(ctx context.Context, query string) (*memory.SearchResponse, error) {
	if c.invocationContext.Memory() == nil {
		return nil, fmt.Errorf("memory service is not set")
	}
	return c.invocationContext.Memory().Search(ctx, query)
}
Repository: google/adk-go
File: cmd/launcher/launcher.go (Go)

Code before:

// SubLauncher is an interface for launchers that can be composed within a parent
// launcher, like the universal launcher. Each SubLauncher corresponds to a
// specific mode of operation (e.g., 'console' or 'web').
type SubLauncher interface {
	// Keyword returns the command-line keyword that activates this sub-launcher.
	Keyword() string
	// Parse parses the arguments for the sub-launcher. It should return any unparsed arguments.
	Parse(args []string) ([]string, error)
	// CommandLineSyntax returns a string describing the command-line flags and arguments for the sub-launcher.
	CommandLineSyntax() string
	// SimpleDescription provides a brief, one-line description of the sub-launcher's function.
	SimpleDescription() string
	// Run executes the sub-launcher's main logic.
	Run(ctx context.Context, config *Config) error
}

// Config contains parameters for web & console execution: sessions, artifacts, agents etc
type Config struct {
	SessionService  session.Service
	ArtifactService artifact.Service
	MemoryService   memory.Service
	AgentLoader     agent.Loader
	A2AOptions      []a2asrv.RequestHandlerOption
}


Review comment (suggestion):
@MarcoCerino23 @baptmont I just noticed, maybe let's make it accept runner.PluginConfig instead of only plugins?

```suggestion
	PluginConfig    runner.PluginConfig
```

Code after:
}

// SubLauncher is an interface for launchers that can be composed within a parent
// launcher, like the universal launcher. Each SubLauncher corresponds to a
// specific mode of operation (e.g., 'console' or 'web').
type SubLauncher interface {
	// Keyword returns the command-line keyword that activates this sub-launcher.
	Keyword() string
	// Parse parses the arguments for the sub-launcher. It should return any unparsed arguments.
	Parse(args []string) ([]string, error)
	// CommandLineSyntax returns a string describing the command-line flags and arguments for the sub-launcher.
	CommandLineSyntax() string
	// SimpleDescription provides a brief, one-line description of the sub-launcher's function.
	SimpleDescription() string
	// Run executes the sub-launcher's main logic.
	Run(ctx context.Context, config *Config) error
}

// Config contains parameters for web & console execution: sessions, artifacts, agents etc
type Config struct {
	SessionService  session.Service
	ArtifactService artifact.Service
	MemoryService   memory.Service
	AgentLoader     agent.Loader
	A2AOptions      []a2asrv.RequestHandlerOption
	Plugins         []*plugin.Plugin
}
Repository: google/comprehensive-rust
File: src/bare-metal/aps/examples/exceptions.S (Assembly)

Code before:
	mrs x1, spsr_el1
	stp x0, x1, [sp, #8 * 22]
.endm

/**
 * Restores the volatile registers from the stack. This currently
 * takes 14 instructions, so it can be used in exception handlers
 * while still leaving 18 instructions left; if paired with
 * save_volatile_to_stack, there are 4 instructions to spare.
 */
.macro restore_volatile_from_stack
	/* Restore registers x2-x18, x29 & x30. */
	ldp x2, x3, [sp, #8 * 2]
	ldp x4, x5, [sp, #8 * 4]
	ldp x6, x7, [sp, #8 * 6]
	ldp x8, x9, [sp, #8 * 8]
	ldp x10, x11, [sp, #8 * 10]
	ldp x12, x13, [sp, #8 * 12]
	ldp x14, x15, [sp, #8 * 14]
	ldp x16, x17, [sp, #8 * 16]
	ldr x18, [sp, #8 * 18]
	ldp x29, x30, [sp, #8 * 20]

	/*
     * Restore registers elr_el1 & spsr_el1, using x0 & x1 as scratch.
     */
	ldp x0, x1, [sp, #8 * 22]
	msr elr_el1, x0
	msr spsr_el1, x1

	/* Restore x0 & x1, and release stack space. */
	ldp x0, x1, [sp], #8 * 24
.endm

/**
 * This is a generic handler for exceptions taken at the current EL
 * while using SP0. It behaves similarly to the SPx case by first
 * switching to SPx, doing the work, then switching back to SP0 before
 * returning.
 *
 * Switching to SPx and calling the Rust handler takes 16
 * instructions. To restore and return we need an additional 16
 * instructions, so we can implement the whole handler within the
 * allotted 32 instructions.
 *
 */
.macro current_exception_sp0 handler:req
	msr spsel, #1
	save_volatile_to_stack
	bl \handler
	restore_volatile_from_stack

Review comment (suggestion):
```suggestion
	/*
	 * Restore registers elr_el1 & spsr_el1, using x0 & x1 as scratch.
	 */
```

Code after:
	mrs x1, spsr_el1
	stp x0, x1, [sp, #8 * 22]
.endm

/**
 * Restores the volatile registers from the stack. This currently
 * takes 14 instructions, so it can be used in exception handlers
 * while still leaving 18 instructions left; if paired with
 * save_volatile_to_stack, there are 4 instructions to spare.
 */
.macro restore_volatile_from_stack
	/* Restore registers x2-x18, x29 & x30. */
	ldp x2, x3, [sp, #8 * 2]
	ldp x4, x5, [sp, #8 * 4]
	ldp x6, x7, [sp, #8 * 6]
	ldp x8, x9, [sp, #8 * 8]
	ldp x10, x11, [sp, #8 * 10]
	ldp x12, x13, [sp, #8 * 12]
	ldp x14, x15, [sp, #8 * 14]
	ldp x16, x17, [sp, #8 * 16]
	ldr x18, [sp, #8 * 18]
	ldp x29, x30, [sp, #8 * 20]

	/*
	 * Restore registers elr_el1 & spsr_el1, using x0 & x1 as scratch.
	 */
	ldp x0, x1, [sp, #8 * 22]
	msr elr_el1, x0
	msr spsr_el1, x1

	/* Restore x0 & x1, and release stack space. */
	ldp x0, x1, [sp], #8 * 24
.endm

/**
 * This is a generic handler for exceptions taken at the current EL
 * while using SP0. It behaves similarly to the SPx case by first
 * switching to SPx, doing the work, then switching back to SP0 before
 * returning.
 *
 * Switching to SPx and calling the Rust handler takes 16
 * instructions. To restore and return we need an additional 16
 * instructions, so we can implement the whole handler within the
 * allotted 32 instructions.
 *
 */
.macro current_exception_sp0 handler:req
	msr spsel, #1
	save_volatile_to_stack
	bl \handler
	restore_volatile_from_stack