AI agents are growing in popularity, leading many of us to connect them with our existing app services to execute key actions. For example, querying existing user data using an AI agent. You can hand the agent a set of tools that have parameters that you would normally define in a non-AI function, something like, fetchUserNotes(userId:string) {}. On the surface, this seems fine to hand off to your AI agent. Having a tool with this method signature can actually lead to a much larger problem since our client app can send whatever they want in the prompt to inject extra information into the request context. In this blog article, we are going to go into detail on what the hidden problem is with this method signature for a tool call.

Tool calling needs extra care

Calling tools may seem straightforward. You may even consider designing your tool calling schemas in the same way you would normally design them in an SDK. However, when it comes to agentic systems, this is not a great idea because it risks exposing more data to the requesting user than is authorized. In the following code snippet, we have a tool definition with two structs and types associated with it.

type GetUserNotesRequest struct {
  UID string `json:"uid" jsonschema:"description=The user id attempting to recover notes"`
}

// this is an anti-pattern, do not do this
getUserNotesTool := genkit.DefineTool(g, 
  "getUserNotes",
  "gets the specified user notes based on the passed in uid",
  func(ctx *ai.ToolContext, input *GetUserNotesRequest) ([]string, error) {
      // ommitted for brevity
})

userInfoFlow := genkit.DefineFlow(g,
  "userInfoFlow",
  func(ctx context.Context, input string) (string, error) {
  ac := core.FromContext(ctx)
  // ommitted decoding app check token

  // ommitted reading UID
  

  mr, err := genkit.Generate(ctx, g,
      ai.WithPrompt(
          fmt.Sprintf(`
          You are a cheerful helpful assistant for the notes app. The user id
          of the currently logged in user is %s. Please help the user complete
          their request : %s`, uid, input
          )
      ),
      ai.WithTools(getUserNotesTool))
      if err != nil {
      return nil, err
      }

  return mr.Text(), nil
})

You can see that the input is a GetUserNotesRequest which takes a user ID as the input value into the function and outputs a list of notes. Can you spot what may be wrong with this definition of a tool?

What’s wrong with the tool definition?

You will notice that in the genkit.Generate call, we supply a prompt that supplies the user id to the AI system via us injecting it with the logged in user. The problem is that as we make that call, our end user can insert a note in their prompt to fetch all user notes for a different user id, i.e., fetch all notes for user ‘1234’. This removes our additional context that we provide in the prompt and returns arbitrary user records for any user when the tool call is made.

Using application context and least privilege

There are two separate and distinct ways to handle this type of design and issue within AI applications. One would be to use the context that your application is operating in and the other is to design a least privileged system for access. Let’s start with the context design first.

Context

In a context designed system, you likely have the user id from the request because you are already restricting requests to authorized users, but what if you wanted to pass that context to other tools? You can use a built-in context object offered in lots of AI platforms to send that additional context in requests for each tool request you want to make. When making requests in Genkit, you parse the headers and extract the Authorization token. You can parse the token and extract the UID as in our example or you could pass the entire token around to different methods. An example of this may look like the following code snippet:

mux := http.NewServeMux()
mux.HandleFunc("POST /userInfoFlow",
  genkit.Handler(userInfoFlow,
      genkit.WithContextProviders(
          func(
              ctx context.Context, req core.RequestData
          ) (core.ActionContext, error) {
              ac := make(map[string]any)

              // get userID from bearer token
              ac["uid"] = ""
              authTokenResult, err := auth.VerifyIDTokenAndCheckRevoked(
                  ctx, strings.TrimPrefix(req.Headers["authorization"], "Bearer "))
              if err == nil {
              ac["uid"] = authTokenResult.UID
              }

              // add app check token to context
              ac["appchecktoken"] = req.Headers["x-firebase-appcheck"]

              return ac, nil
})))

This context can then be used in other locations. For instance, looking back at our getUserNotesTool the tool definition may now look like the following.

type EmptyToolRequest struct{}

getUserNotesTool := genkit.DefineTool(g,
  "getUserNotes",
  "gets the specified user notes based on the passed in uid",
  func(ctx *ai.ToolContext, input *EmptyToolRequest) (string, error) {
    // the uid is supplied from the Headers
    uid := core.FromContext(ctx.Context)["uid"]
    v := fetchNotesFromDb(uid)
    return v, nil
})

Now, instead of fetching whatever user ID was passed to it, the getUserNotes function will only fetch notes for the uid that was retrieved from the Authorization token.

Least privilege

Even better, make it impossible for the getUserNotes tool to read another user’s notes. In our example, we were using an account with a heightened privilege to the database. This is great for certain scenarios where we need heightened access like correlating data between multiple tables and accounts. However, the AI workflow we created is user centric and should focus on the access that the user has, not what an administrator has access to.

Our getUserNotes tool should really only be able to access the notes that the user has access to, not all notes. So to make this better, we should instead replace our heightened account request with a user level request. This means we should not call our backend with a highly privileged account that can access all tables. Instead, we should call our backend with the user’s credential on a user-available endpoint. This limits the amount of data that can be accessed with that API key. This would provide an extra layer of security in the event the user found another way to switch the tool call to calling a different user id in our example.

Summary

Securing AI agents is more than just securing the authentication and authorization to those agents, it even comes down to the implementation of the tools that the agents have access to. We can start by passing in the appropriate context to each tool via a context object so the AI is not left to infer the appropriate context. Finally, we can also use the lower privileged SDKs to communicate on behalf of the user rather than some of the higher privilege admin SDKs so we can ensure that users do not accidentally get more access than they normally would have.

package main

import (
"context"
"fmt"
"log"
"net/http"
"os"
"os/signal"
"strings"
"time"

firebase "firebase.google.com/go/v4"
"github.com/firebase/genkit/go/ai"
"github.com/firebase/genkit/go/core"
"github.com/firebase/genkit/go/genkit"
"github.com/firebase/genkit/go/plugins/googlegenai"
)

const (
PROJECT_ID     string = "PROJECT_ID"
CLOUD_LOCATION string = "REGION_STRING"
)

func main() {
ctx := context.Background()

app, err := firebase.NewApp(ctx, &firebase.Config{ProjectID: PROJECT_ID})
if err != nil {
	log.Fatalf("error connecting to firebase, cannot run server : %v
", err)
}

auth, err := app.Auth(ctx)
if err != nil {
	log.Fatalf("cannot connect to auth %v
", err)
}

appCheck, err := app.AppCheck(ctx)
if err != nil {
	log.Fatalf("cannot connect to app check %v
", err)
}

g := genkit.Init(ctx,
	genkit.WithPlugins(&googlegenai.VertexAI{ProjectID: PROJECT_ID, Location: CLOUD_LOCATION}),
	genkit.WithDefaultModel("vertexai/gemini-2.5-flash"),
)

type GetUserNotesRequest struct {
	UID string `json:"uid" jsonschema:"description=The user id attempting to recover notes"`
}

type GetUserNotesResponse struct {
	Notes []string `json:"notes" jsonschema:"description=Each note the user has recorded in the app"`
}

type EmptyToolRequest struct{}

getUserNotesTool := genkit.DefineTool(g, "getUserNotes", "gets all of the logged in users notes",
	func(ctx *ai.ToolContext, input *EmptyToolRequest) (string, error) {
		uid := core.FromContext(ctx.Context)["uid"]
		v := fmt.Sprintf("The user requesting data is %s, but the user requested data is for user %s", uid.(string), uid.(string))
		fmt.Println(v)
		return v, nil
	})

type UserInfoRequest struct {
	Request string `json:"request" jsonschema:"description=The user request for a specific task to help with their notes app"`
}

type UserInfoResponse struct {
	Response string `json:"response" jsonschema:"description=The helpful response from the AI completing the user request"`
}
userInfoFlow := genkit.DefineFlow(g, "userInfoFlow", func(ctx context.Context, input *UserInfoRequest) (*UserInfoResponse, error) {
	ac := core.FromContext(ctx)
	_, err := appCheck.VerifyToken(ac["appchecktoken"].(string))
	if err != nil {
		return nil, fmt.Errorf("could not decode a valid app check token. attestation likely failed")
	}
	uid := ac["uid"]
	if uid == "" {
		return nil, fmt.Errorf("could not get a user id. are you authenticated?")
	}

	mr, err := genkit.Generate(ctx, g, ai.WithPrompt(fmt.Sprintf("You are a cheerful helpful assistant of an AI. Please help the user complete their request : %s", input.Request)), ai.WithTools(getUserNotesTool))
	if err != nil {
		return nil, err
	}
	response := mr.Text()

	return &UserInfoResponse{Response: response}, nil
})

mux := http.NewServeMux()
mux.HandleFunc("POST /userInfoFlow", genkit.Handler(userInfoFlow, genkit.WithContextProviders(func(ctx context.Context, req core.RequestData) (core.ActionContext, error) {
	ac := make(map[string]any)
	ac["uid"] = ""
	ac["appchecktoken"] = ""

	authTokenResult, err := auth.VerifyIDTokenAndCheckRevoked(ctx, strings.TrimPrefix(req.Headers["authorization"], "Bearer "))
	if err == nil {
		ac["uid"] = authTokenResult.UID
	}

	ac["appchecktoken"] = req.Headers["x-firebase-appcheck"]

	return ac, nil
})))
server := &http.Server{
	Addr:    "127.0.0.1:3400",
	Handler: corsMiddleware(mux),
}
go func() {
	log.Println("Starting server on http://localhost:3400")
	log.Println("Flow available at: POST http://localhost:3400/userInfoFlow")
	if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
		log.Fatalf("ListenAndServe(): %v", err)
	}
}()

quit := make(chan os.Signal, 1)
signal.Notify(quit, os.Interrupt)
<-quit
log.Println("Shutting down server...")

ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()

if err := server.Shutdown(ctx); err != nil {
	log.Fatalf("Server Shutdown Failed:%+v", err)
}
log.Println("Server exited properly")
}

func corsMiddleware(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Access-Control-Allow-Origin", "https://your-web-app.com")
	w.Header().Set("Access-Control-Allow-Methods", "POST, GET, OPTIONS, PUT, DELETE")
	w.Header().Set("Access-Control-Allow-Headers", "Accept, Content-Type, Content-Length, Accept-Encoding, X-CSRF-Token, Authorization, X-Firebase-AppCheck")
	if r.Method == "OPTIONS" {
		w.WriteHeader(http.StatusOK)
		return
	}
	next.ServeHTTP(w, r)
})
}