feat: add new user-agent-extra flag

## Description

Addresses an issue where Dataplex AspectTypes created during integration
tests were not consistently deleted. This accumulation led to exceeding
the AspectType quota.
To prevent this, the test setup now includes a step to list and delete
all existing AspectTypes within the test project and location *before*
attempting to create any new ones. This ensures a clean state for each
test run and avoids hitting the quota.


## PR Checklist

> Thank you for opening a Pull Request! Before submitting your PR, there
are a
> few things you can do to make sure it goes smoothly:

- [ ] Make sure you reviewed

[CONTRIBUTING.md](https://github.com/googleapis/genai-toolbox/blob/main/CONTRIBUTING.md)
- [ ] Make sure to open an issue as a

[bug/issue](https://github.com/googleapis/genai-toolbox/issues/new/choose)
  before writing your code! That way we can discuss the change, evaluate
  designs, and agree on the general idea
- [ ] Ensure the tests and linter pass
- [ ] Code coverage does not decrease (if any source code was changed)
- [ ] Appropriate docs were updated (if necessary)
- [ ] Make sure to add `!` if this involve a breaking change

🛠️ Fixes #2057

Co-authored-by: Averi Kitsch <akitsch@google.com>

cmd/root.go

@@ -384,6 +384,7 @@ func NewCommand(opts ...Option) *Command {
 	// TODO: Insecure by default. Might consider updating this for v1.0.0
 	flags.StringSliceVar(&cmd.cfg.AllowedOrigins, "allowed-origins", []string{"*"}, "Specifies a list of origins permitted to access this server. Defaults to '*'.")
 	flags.StringSliceVar(&cmd.cfg.AllowedHosts, "allowed-hosts", []string{"*"}, "Specifies a list of hosts permitted to access this server. Defaults to '*'.")
+	flags.StringSliceVar(&cmd.cfg.UserAgentExtra, "user-agent-extra", []string{}, "Appends additional metadata to the User-Agent.")
 
 	// wrap RunE command so that we have access to original Command object
 	cmd.RunE = func(*cobra.Command, []string) error { return run(cmd) }

cmd/root_test.go

@@ -70,6 +70,9 @@ func withDefaults(c server.ServerConfig) server.ServerConfig {
 	if c.AllowedHosts == nil {
 		c.AllowedHosts = []string{"*"}
 	}
+	if c.UserAgentExtra == nil {
+		c.UserAgentExtra = []string{}
+	}
 	return c
 }
 
@@ -230,6 +233,13 @@ func TestServerConfigFlags(t *testing.T) {
 				AllowedHosts: []string{"http://foo.com", "http://bar.com"},
 			}),
 		},
+		{
+			desc: "user agent extra",
+			args: []string{"--user-agent-extra", "foo,bar"},
+			want: withDefaults(server.ServerConfig{
+				UserAgentExtra: []string{"foo", "bar"},
+			}),
+		},
 	}
 	for _, tc := range tcs {
 		t.Run(tc.desc, func(t *testing.T) {

docs/en/getting-started/colab_quickstart.ipynb

@@ -509,7 +509,7 @@
       },
       "outputs": [],
       "source": [
-        "! pip install toolbox-adk --quiet\n",
+        "! pip install toolbox-core --quiet\n",
         "! pip install google-adk --quiet"
       ]
     },
@@ -525,18 +525,14 @@
         "from google.adk.runners import Runner\n",
         "from google.adk.sessions import InMemorySessionService\n",
         "from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService\n",
-        "from google.adk.tools.toolbox_toolset import ToolboxToolset\n",
         "from google.genai import types\n",
+        "from toolbox_core import ToolboxSyncClient\n",
         "\n",
         "import os\n",
         "# TODO(developer): replace this with your Google API key\n",
         "os.environ['GOOGLE_API_KEY'] = \"<GOOGLE_API_KEY>\"\n",
         "\n",
-        "# Configure toolset\n",
-        "toolset = ToolboxToolset(\n",
-        "    server_url=\"http://127.0.0.1:5000\",\n",
-        "    toolset_name=\"my-toolset\"\n",
-        ")\n",
+        "toolbox_client = ToolboxSyncClient(\"http://127.0.0.1:5000\")\n",
         "\n",
         "prompt = \"\"\"\n",
         "  You're a helpful hotel assistant. You handle hotel searching, booking and\n",
@@ -553,7 +549,7 @@
         "    name='hotel_agent',\n",
         "    description='A helpful AI assistant.',\n",
         "    instruction=prompt,\n",
-        "    tools=[toolset],\n",
+        "    tools=toolbox_client.load_toolset(\"my-toolset\"),\n",
         ")\n",
         "\n",
         "session_service = InMemorySessionService()\n",

docs/en/getting-started/local_quickstart.md

@@ -52,7 +52,7 @@ runtime](https://research.google.com/colaboratory/local-runtimes.html).
     {{< tabpane persist=header >}}
 {{< tab header="ADK" lang="bash" >}}
 
-pip install toolbox-adk
+pip install toolbox-core
 {{< /tab >}}
 {{< tab header="Langchain" lang="bash" >}}
 

docs/en/getting-started/quickstart/python/adk/quickstart.py

@@ -1,17 +1,15 @@
 from google.adk import Agent
 from google.adk.apps import App
-from google.adk.tools.toolbox_toolset import ToolboxToolset
+from toolbox_core import ToolboxSyncClient
 
 # TODO(developer): update the TOOLBOX_URL to your toolbox endpoint
-toolset = ToolboxToolset(
-    server_url="http://127.0.0.1:5000",
-)
+client = ToolboxSyncClient("http://127.0.0.1:5000")
 
 root_agent = Agent(
     name='root_agent',
     model='gemini-2.5-flash',
     instruction="You are a helpful AI assistant designed to provide accurate and useful information.",
-    tools=[toolset],
+    tools=client.load_toolset(),
 )
 
 app = App(root_agent=root_agent, name="my_agent")

docs/en/getting-started/quickstart/python/adk/requirements.txt

@@ -1,3 +1,3 @@
 google-adk==1.21.0
-toolbox-adk>=0.1.0
+toolbox-core==0.5.4
 pytest==9.0.2

docs/en/how-to/deploy_adk_agent.md

@@ -49,7 +49,7 @@ with the necessary configuration for deployment to Vertex AI Agent Engine.
 4.  Add `toolbox-core` as a dependency to the new project:
 
     ```bash
-    uv add toolbox-adk
+    uv add toolbox-core
     ```
 
 ## Step 3: Configure Google Cloud Authentication
@@ -95,23 +95,22 @@ authentication token.
     ```python
     from google.adk import Agent
     from google.adk.apps import App
-    from google.adk.tools.toolbox_toolset import ToolboxToolset
-    from toolbox_adk import CredentialStrategy
+    from toolbox_core import ToolboxSyncClient, auth_methods
 
     # TODO(developer): Replace with your Toolbox Cloud Run Service URL
     TOOLBOX_URL = "https://your-toolbox-service-xyz.a.run.app"
 
-    # Initialize the toolset with Workload Identity (generates ID token for the URL)
-    toolset = ToolboxToolset(
-        server_url=TOOLBOX_URL,
-        credentials=CredentialStrategy.workload_identity(target_audience=TOOLBOX_URL)
+    # Initialize the client with the Cloud Run URL and Auth headers
+    client = ToolboxSyncClient(
+        TOOLBOX_URL, 
+        client_headers={"Authorization": auth_methods.get_google_id_token(TOOLBOX_URL)}
     )
 
     root_agent = Agent(
         name='root_agent',
         model='gemini-2.5-flash',
         instruction="You are a helpful AI assistant designed to provide accurate and useful information.",
-        tools=[toolset],
+        tools=client.load_toolset(),
     )
 
     app = App(root_agent=root_agent, name="my_agent")

docs/en/reference/cli.md

@@ -27,6 +27,7 @@ description: >
 |              | `--ui`                     | Launches the Toolbox UI web server.                                                                                                                                              |             |
 |              | `--allowed-origins`        | Specifies a list of origins permitted to access this server for CORs access.                                                                                                     | `*`         |
 |              | `--allowed-hosts`          | Specifies a list of hosts permitted to access this server to prevent DNS rebinding attacks.                                                                                      | `*`         |
+|              | `--user-agent-extra`       | Appends additional metadata to the User-Agent.                                                                                                                                   |             |
 | `-v`         | `--version`                | version for toolbox                                                                                                                                                              |             |
 
 ## Examples

docs/en/samples/bigquery/local_quickstart.md

@@ -365,7 +365,7 @@ pip install llama-index-llms-google-genai
 
 {{< /tab >}}
 {{< tab header="ADK" lang="bash" >}}
-pip install toolbox-adk
+pip install toolbox-core
 {{< /tab >}}
 {{< /tabpane >}}
 
@@ -607,8 +607,8 @@ from google.adk.agents import Agent
 from google.adk.runners import Runner
 from google.adk.sessions import InMemorySessionService
 from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService
-from google.adk.tools.toolbox_toolset import ToolboxToolset
 from google.genai import types # For constructing message content
+from toolbox_core import ToolboxSyncClient
 
 import os
 os.environ['GOOGLE_GENAI_USE_VERTEXAI'] = 'True'
@@ -623,47 +623,48 @@ os.environ['GOOGLE_CLOUD_LOCATION'] = 'us-central1'
 
 # --- Load Tools from Toolbox ---
 
-# TODO(developer): Ensure the Toolbox server is running at http://127.0.0.1:5000
-toolset = ToolboxToolset(server_url="http://127.0.0.1:5000")
+# TODO(developer): Ensure the Toolbox server is running at <http://127.0.0.1:5000>
 
-# --- Define the Agent's Prompt ---
-prompt = """
-  You're a helpful hotel assistant. You handle hotel searching, booking and
-  cancellations. When the user searches for a hotel, mention it's name, id,
-  location and price tier. Always mention hotel ids while performing any
-  searches. This is very important for any operations. For any bookings or
-  cancellations, please provide the appropriate confirmation. Be sure to
-  update checkin or checkout dates if mentioned by the user.
-  Don't ask for confirmations from the user.
-"""
+with ToolboxSyncClient("<http://127.0.0.1:5000>") as toolbox_client:
+    # TODO(developer): Replace "my-toolset" with the actual ID of your toolset as configured in your MCP Toolbox server.
+    agent_toolset = toolbox_client.load_toolset("my-toolset")
 
-# --- Configure the Agent ---
+    # --- Define the Agent's Prompt ---
+    prompt = """
+      You're a helpful hotel assistant. You handle hotel searching, booking and
+      cancellations. When the user searches for a hotel, mention it's name, id,
+      location and price tier. Always mention hotel ids while performing any
+      searches. This is very important for any operations. For any bookings or
+      cancellations, please provide the appropriate confirmation. Be sure to
+      update checkin or checkout dates if mentioned by the user.
+      Don't ask for confirmations from the user.
+    """
 
-root_agent = Agent(
-    model='gemini-2.0-flash-001',
-    name='hotel_agent',
-    description='A helpful AI assistant that can search and book hotels.',
-    instruction=prompt,
-    tools=[toolset], # Pass the loaded toolset
-)
+    # --- Configure the Agent ---
 
-# --- Initialize Services for Running the Agent ---
-session_service = InMemorySessionService()
-artifacts_service = InMemoryArtifactService()
+    root_agent = Agent(
+        model='gemini-2.0-flash-001',
+        name='hotel_agent',
+        description='A helpful AI assistant that can search and book hotels.',
+        instruction=prompt,
+        tools=agent_toolset, # Pass the loaded toolset
+    )
 
-runner = Runner(
-    app_name='hotel_agent',
-    agent=root_agent,
-    artifact_service=artifacts_service,
-    session_service=session_service,
-)
-
-async def main():
+    # --- Initialize Services for Running the Agent ---
+    session_service = InMemorySessionService()
+    artifacts_service = InMemoryArtifactService()
     # Create a new session for the interaction.
-    session = await session_service.create_session(
+    session = session_service.create_session(
         state={}, app_name='hotel_agent', user_id='123'
     )
 
+    runner = Runner(
+        app_name='hotel_agent',
+        agent=root_agent,
+        artifact_service=artifacts_service,
+        session_service=session_service,
+    )
+
     # --- Define Queries and Run the Agent ---
     queries = [
         "Find hotels in Basel with Basel in it's name.",
@@ -686,10 +687,6 @@ async def main():
 
         for text in responses:
           print(text)
-
-import asyncio
-if __name__ == "__main__":
-    asyncio.run(main())
 {{< /tab >}}
 {{< /tabpane >}}
 

internal/server/config.go

@@ -64,12 +64,14 @@ type ServerConfig struct {
 	Stdio bool
 	// DisableReload indicates if the user has disabled dynamic reloading for Toolbox.
 	DisableReload bool
-	// UI indicates if Toolbox UI endpoints (/ui) are available
+	// UI indicates if Toolbox UI endpoints (/ui) are available.
 	UI bool
 	// Specifies a list of origins permitted to access this server.
 	AllowedOrigins []string
-	// Specifies a list of hosts permitted to access this server
+	// Specifies a list of hosts permitted to access this server.
 	AllowedHosts []string
+	// UserAgentExtra specifies additional metadata to append to the User-Agent string.
+	UserAgentExtra []string
 }
 
 type logFormat string

internal/server/server.go

@@ -64,7 +64,11 @@ func InitializeConfigs(ctx context.Context, cfg ServerConfig) (
 	map[string]prompts.Promptset,
 	error,
 ) {
-	ctx = util.WithUserAgent(ctx, cfg.Version)
+	metadataStr := cfg.Version
+	if len(cfg.UserAgentExtra) > 0 {
+		metadataStr += "+" + strings.Join(cfg.UserAgentExtra, ".")
+	}
+	ctx = util.WithUserAgent(ctx, metadataStr)
 	instrumentation, err := util.InstrumentationFromContext(ctx)
 	if err != nil {
 		panic(err)

tests/dataplex/dataplex_integration_test.go

@@ -85,13 +85,66 @@ func initDataplexConnection(ctx context.Context) (*dataplex.CatalogClient, error
 	return client, nil
 }
 
+// cleanupOldAspectTypes Deletes AspectTypes older than the specified duration.
+func cleanupOldAspectTypes(t *testing.T, ctx context.Context, client *dataplex.CatalogClient, oldThreshold time.Duration) {
+	parent := fmt.Sprintf("projects/%s/locations/us", DataplexProject)
+	olderThanTime := time.Now().Add(-oldThreshold)
+
+	listReq := &dataplexpb.ListAspectTypesRequest{
+		Parent:   parent,
+		PageSize: 100,               // Fetch up to 100 items
+		OrderBy:  "create_time asc", // Order by creation time
+	}
+
+	const maxDeletes = 8 // Explicitly limit the number of deletions
+	it := client.ListAspectTypes(ctx, listReq)
+	var aspectTypesToDelete []string
+	for len(aspectTypesToDelete) < maxDeletes {
+		aspectType, err := it.Next()
+		if err == iterator.Done {
+			break
+		}
+		if err != nil {
+			t.Logf("Warning: Failed to list aspect types during cleanup: %v", err)
+			return
+		}
+		// Perform time-based filtering in memory
+		if aspectType.CreateTime != nil {
+			createTime := aspectType.CreateTime.AsTime()
+			if createTime.Before(olderThanTime) {
+				aspectTypesToDelete = append(aspectTypesToDelete, aspectType.GetName())
+			}
+		} else {
+			t.Logf("Warning: AspectType %s has no CreateTime", aspectType.GetName())
+		}
+	}
+	if len(aspectTypesToDelete) == 0 {
+		t.Logf("cleanupOldAspectTypes: No aspect types found older than %s to delete.", oldThreshold.String())
+		return
+	}
+
+	for _, aspectTypeName := range aspectTypesToDelete {
+		deleteReq := &dataplexpb.DeleteAspectTypeRequest{Name: aspectTypeName}
+		op, err := client.DeleteAspectType(ctx, deleteReq)
+		if err != nil {
+			t.Logf("Warning: Failed to delete aspect type %s: %v", aspectTypeName, err)
+			continue // Skip to the next item if initiation fails
+		}
+
+		if err := op.Wait(ctx); err != nil {
+			t.Logf("Warning: Failed to delete aspect type %s, operation error: %v", aspectTypeName, err)
+		} else {
+			t.Logf("cleanupOldAspectTypes: Successfully deleted %s", aspectTypeName)
+		}
+	}
+}
+
 func TestDataplexToolEndpoints(t *testing.T) {
 	sourceConfig := getDataplexVars(t)
 	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Minute)
 	defer cancel()
 
 	var args []string
-
 	bigqueryClient, err := initBigQueryConnection(ctx, DataplexProject)
 	if err != nil {
 		t.Fatalf("unable to create Cloud SQL connection pool: %s", err)
@@ -102,6 +155,9 @@ func TestDataplexToolEndpoints(t *testing.T) {
 		t.Fatalf("unable to create Dataplex connection: %s", err)
 	}
 
+	// Cleanup older aspecttypes
+	cleanupOldAspectTypes(t, ctx, dataplexClient, 1*time.Hour)
+
 	// create resources with UUID
 	datasetName := fmt.Sprintf("temp_toolbox_test_%s", strings.ReplaceAll(uuid.New().String(), "-", ""))
 	tableName := fmt.Sprintf("param_table_%s", strings.ReplaceAll(uuid.New().String(), "-", ""))