Add API_KEY to secretEnv in cloudbuild.yaml

Added API_KEY to secret environment variables for integration.

.ci/integration.cloudbuild.yaml

@@ -87,7 +87,7 @@ steps:
       - "CLOUD_SQL_POSTGRES_REGION=$_REGION"
       - "SERVICE_ACCOUNT_EMAIL=$SERVICE_ACCOUNT_EMAIL"
     secretEnv:
-      ["CLOUD_SQL_POSTGRES_USER", "CLOUD_SQL_POSTGRES_PASS", "CLIENT_ID"]
+      ["CLOUD_SQL_POSTGRES_USER", "CLOUD_SQL_POSTGRES_PASS", "CLIENT_ID", "API_KEY"]
     volumes:
       - name: "go"
         path: "/gopath"
@@ -134,7 +134,7 @@ steps:
       - "ALLOYDB_POSTGRES_DATABASE=$_DATABASE_NAME"
       - "ALLOYDB_POSTGRES_REGION=$_REGION"
       - "SERVICE_ACCOUNT_EMAIL=$SERVICE_ACCOUNT_EMAIL"
-    secretEnv: ["ALLOYDB_POSTGRES_USER", "ALLOYDB_POSTGRES_PASS", "CLIENT_ID"]
+    secretEnv: ["ALLOYDB_POSTGRES_USER", "ALLOYDB_POSTGRES_PASS", "CLIENT_ID", "API_KEY"]
     volumes:
       - name: "go"
         path: "/gopath"
@@ -305,7 +305,7 @@ steps:
       - "POSTGRES_HOST=$_POSTGRES_HOST"
       - "POSTGRES_PORT=$_POSTGRES_PORT"
       - "SERVICE_ACCOUNT_EMAIL=$SERVICE_ACCOUNT_EMAIL"
-    secretEnv: ["POSTGRES_USER", "POSTGRES_PASS", "CLIENT_ID"]
+    secretEnv: ["POSTGRES_USER", "POSTGRES_PASS", "CLIENT_ID", "API_KEY"]
     volumes:
       - name: "go"
         path: "/gopath"
@@ -964,6 +964,13 @@ steps:
 
 availableSecrets:
   secretManager:
+    # Common secrets
+    - versionName: projects/$PROJECT_ID/secrets/client_id/versions/latest
+      env: CLIENT_ID
+    - versionName: projects/$PROJECT_ID/secrets/api_key/versions/latest
+      env: API_KEY
+
+    # Resource-specific secrets
     - versionName: projects/$PROJECT_ID/secrets/cloud_sql_pg_user/versions/latest
       env: CLOUD_SQL_POSTGRES_USER
     - versionName: projects/$PROJECT_ID/secrets/cloud_sql_pg_pass/versions/latest
@@ -980,8 +987,6 @@ availableSecrets:
       env: POSTGRES_USER
     - versionName: projects/$PROJECT_ID/secrets/postgres_pass/versions/latest
       env: POSTGRES_PASS
-    - versionName: projects/$PROJECT_ID/secrets/client_id/versions/latest
-      env: CLIENT_ID
     - versionName: projects/$PROJECT_ID/secrets/neo4j_user/versions/latest
       env: NEO4J_USER
     - versionName: projects/$PROJECT_ID/secrets/neo4j_pass/versions/latest

cmd/root.go

@@ -98,7 +98,6 @@ import (
 	_ "github.com/googleapis/genai-toolbox/internal/tools/cloudsql/cloudsqlgetinstances"
 	_ "github.com/googleapis/genai-toolbox/internal/tools/cloudsql/cloudsqllistdatabases"
 	_ "github.com/googleapis/genai-toolbox/internal/tools/cloudsql/cloudsqllistinstances"
-	_ "github.com/googleapis/genai-toolbox/internal/tools/cloudsql/cloudsqlrestorebackup"
 	_ "github.com/googleapis/genai-toolbox/internal/tools/cloudsql/cloudsqlwaitforoperation"
 	_ "github.com/googleapis/genai-toolbox/internal/tools/cloudsqlmssql/cloudsqlmssqlcreateinstance"
 	_ "github.com/googleapis/genai-toolbox/internal/tools/cloudsqlmysql/cloudsqlmysqlcreateinstance"

cmd/root_test.go

@@ -1493,7 +1493,7 @@ func TestPrebuiltTools(t *testing.T) {
 			wantToolset: server.ToolsetConfigs{
 				"cloud_sql_postgres_admin_tools": tools.ToolsetConfig{
 					Name:      "cloud_sql_postgres_admin_tools",
-					ToolNames: []string{"create_instance", "get_instance", "list_instances", "create_database", "list_databases", "create_user", "wait_for_operation", "postgres_upgrade_precheck", "clone_instance", "create_backup", "restore_backup"},
+					ToolNames: []string{"create_instance", "get_instance", "list_instances", "create_database", "list_databases", "create_user", "wait_for_operation", "postgres_upgrade_precheck", "clone_instance", "create_backup"},
 				},
 			},
 		},
@@ -1503,7 +1503,7 @@ func TestPrebuiltTools(t *testing.T) {
 			wantToolset: server.ToolsetConfigs{
 				"cloud_sql_mysql_admin_tools": tools.ToolsetConfig{
 					Name:      "cloud_sql_mysql_admin_tools",
-					ToolNames: []string{"create_instance", "get_instance", "list_instances", "create_database", "list_databases", "create_user", "wait_for_operation", "clone_instance", "create_backup", "restore_backup"},
+					ToolNames: []string{"create_instance", "get_instance", "list_instances", "create_database", "list_databases", "create_user", "wait_for_operation", "clone_instance", "create_backup"},
 				},
 			},
 		},
@@ -1513,7 +1513,7 @@ func TestPrebuiltTools(t *testing.T) {
 			wantToolset: server.ToolsetConfigs{
 				"cloud_sql_mssql_admin_tools": tools.ToolsetConfig{
 					Name:      "cloud_sql_mssql_admin_tools",
-					ToolNames: []string{"create_instance", "get_instance", "list_instances", "create_database", "list_databases", "create_user", "wait_for_operation", "clone_instance", "create_backup", "restore_backup"},
+					ToolNames: []string{"create_instance", "get_instance", "list_instances", "create_database", "list_databases", "create_user", "wait_for_operation", "clone_instance", "create_backup"},
 				},
 			},
 		},

docs/en/how-to/connect-ide/cloud_sql_mssql_admin_mcp.md

@@ -54,7 +54,6 @@ instance, database and users:
         * `create_instance`
         * `create_user`
         * `clone_instance`
-        * `restore_backup`
 
 ## Install MCP Toolbox
 
@@ -302,7 +301,6 @@ instances and interacting with your database:
 * **wait_for_operation**: Waits for a Cloud SQL operation to complete.
 * **clone_instance**: Creates a clone of an existing Cloud SQL for SQL Server instance.
 * **create_backup**: Creates a backup on a Cloud SQL instance.
-* **restore_backup**: Restores a backup of a Cloud SQL instance.
 
 {{< notice note >}}
 Prebuilt tools are pre-1.0, so expect some tool changes between versions. LLMs

docs/en/how-to/connect-ide/cloud_sql_mysql_admin_mcp.md

@@ -54,7 +54,6 @@ database and users:
         * `create_instance`
         * `create_user`
         * `clone_instance`
-        * `restore_backup`
 
 ## Install MCP Toolbox
 
@@ -302,7 +301,6 @@ instances and interacting with your database:
 * **wait_for_operation**: Waits for a Cloud SQL operation to complete.
 * **clone_instance**: Creates a clone of an existing Cloud SQL for MySQL instance.
 * **create_backup**: Creates a backup on a Cloud SQL instance.
-* **restore_backup**: Restores a backup of a Cloud SQL instance.
 
 {{< notice note >}}
 Prebuilt tools are pre-1.0, so expect some tool changes between versions. LLMs

docs/en/how-to/connect-ide/cloud_sql_pg_admin_mcp.md

@@ -54,7 +54,6 @@ instance, database and users:
         * `create_instance`
         * `create_user`
         * `clone_instance`
-        * `restore_backup`
 
 ## Install MCP Toolbox
 
@@ -302,7 +301,6 @@ instances and interacting with your database:
 * **wait_for_operation**: Waits for a Cloud SQL operation to complete.
 * **clone_instance**: Creates a clone of an existing Cloud SQL for PostgreSQL instance.
 * **create_backup**: Creates a backup on a Cloud SQL instance.
-* **restore_backup**: Restores a backup of a Cloud SQL instance.
 
 {{< notice note >}}
 Prebuilt tools are pre-1.0, so expect some tool changes between versions. LLMs

docs/en/reference/prebuilt-tools.md

@@ -194,7 +194,6 @@ See [Usage Examples](../reference/cli.md#examples).
         * `create_instance`
         * `create_user`
         * `clone_instance`
-        * `restore_backup`
 
 *   **Tools:**
     *   `create_instance`: Creates a new Cloud SQL for MySQL instance.
@@ -206,7 +205,6 @@ See [Usage Examples](../reference/cli.md#examples).
     *   `wait_for_operation`: Waits for a Cloud SQL operation to complete.
     *   `clone_instance`: Creates a clone for an existing Cloud SQL for MySQL instance.
     *   `create_backup`: Creates a backup on a Cloud SQL instance.
-    *   `restore_backup`: Restores a backup of a Cloud SQL instance.
 
 ## Cloud SQL for PostgreSQL
 
@@ -286,7 +284,6 @@ See [Usage Examples](../reference/cli.md#examples).
         * `create_instance`
         * `create_user`
         * `clone_instance`
-        * `restore_backup`
 *   **Tools:**
     *   `create_instance`: Creates a new Cloud SQL for PostgreSQL instance.
     *   `get_instance`: Gets information about a Cloud SQL instance.
@@ -297,7 +294,6 @@ See [Usage Examples](../reference/cli.md#examples).
     *   `wait_for_operation`: Waits for a Cloud SQL operation to complete.
     *   `clone_instance`: Creates a clone for an existing Cloud SQL for PostgreSQL instance.
     *   `create_backup`: Creates a backup on a Cloud SQL instance.
-    *   `restore_backup`: Restores a backup of a Cloud SQL instance.
 
 ## Cloud SQL for SQL Server
 
@@ -351,7 +347,6 @@ See [Usage Examples](../reference/cli.md#examples).
         * `create_instance`
         * `create_user`
         * `clone_instance`
-        * `restore_backup`
 *   **Tools:**
     *   `create_instance`: Creates a new Cloud SQL for SQL Server instance.
     *   `get_instance`: Gets information about a Cloud SQL instance.
@@ -362,7 +357,6 @@ See [Usage Examples](../reference/cli.md#examples).
     *   `wait_for_operation`: Waits for a Cloud SQL operation to complete.
     *   `clone_instance`: Creates a clone for an existing Cloud SQL for SQL Server instance.
     *   `create_backup`: Creates a backup on a Cloud SQL instance.
-    *   `restore_backup`: Restores a backup of a Cloud SQL instance.
 
 ## Dataplex
 

docs/en/resources/sources/dgraph.md

@@ -7,17 +7,6 @@ description: >
 
 ---
 
-{{< notice note >}}
-**⚠️ Best Effort Maintenance**
-
-This integration is maintained on a best-effort basis by the project
-team/community. While we strive to address issues and provide workarounds when
-resources are available, there are no guaranteed response times or code fixes.
-
-The automated integration tests for this module are currently non-functional or
-failing.
-{{< /notice >}}
-
 ## About
 
 [Dgraph][dgraph-docs] is an open-source graph database. It is designed for

docs/en/resources/tools/cloudsql/cloudsqlrestorebackup.md

@@ -1,53 +0,0 @@
----
-title: cloud-sql-restore-backup
-type: docs
-weight: 10
-description: "Restores a backup of a Cloud SQL instance."
----
-
-The `cloud-sql-restore-backup` tool restores a backup on a Cloud SQL instance using the Cloud SQL Admin API.
-
-{{< notice info dd>}}
-This tool uses a `source` of kind `cloud-sql-admin`.
-{{< /notice >}}
-
-## Examples
-
-Basic backup restore
-
-```yaml
-tools:
-  backup-restore-basic:
-    kind: cloud-sql-restore-backup
-    source: cloud-sql-admin-source
-    description: "Restores a backup onto the given Cloud SQL instance."
-```
-
-## Reference
-
-### Tool Configuration
-| **field**      | **type** | **required** | **description**                                  |
-| -------------- | :------: | :----------: | ------------------------------------------------ |
-| kind           | string   | true         | Must be "cloud-sql-restore-backup".              |
-| source         | string   | true         | The name of the `cloud-sql-admin` source to use. |
-| description    | string   | false        | A description of the tool.                       |
-
-### Tool Inputs
-
-| **parameter**     | **type** | **required** | **description**                                                              |
-| ------------------| :------: | :----------: | -----------------------------------------------------------------------------|
-| target_project    | string   | true         | The project ID of the instance to restore the backup onto.                   |
-| target_instance   | string   | true         | The instance to restore the backup onto. Does not include the project ID.    |
-| backup_id         | string   | true         | The identifier of the backup being restored.                                 |
-| source_project    | string   | false        | (Optional) The project ID of the instance that the backup belongs to.        |
-| source_instance   | string   | false        | (Optional) Cloud SQL instance ID of the instance that the backup belongs to. |
-
-## Usage Notes
-
-- The `backup_id` field can be a BackupRun ID (which will be an int64), backup name, or BackupDR backup name.
-- If the `backup_id` field contains a BackupRun ID (i.e. an int64), the optional fields `source_project` and `source_instance` must also be provided.
-
-## See Also
-- [Cloud SQL Admin API documentation](https://cloud.google.com/sql/docs/mysql/admin-api)
-- [Toolbox Cloud SQL tools documentation](../cloudsql)
-- [Cloud SQL Restore API documentation](https://cloud.google.com/sql/docs/mysql/backup-recovery/restoring)

docs/en/resources/tools/dgraph/dgraph-dql.md

@@ -9,17 +9,6 @@ aliases:
 - /resources/tools/dgraph-dql
 ---
 
-{{< notice note >}}
-**⚠️ Best Effort Maintenance**
-
-This integration is maintained on a best-effort basis by the project
-team/community. While we strive to address issues and provide workarounds when
-resources are available, there are no guaranteed response times or code fixes.
-
-The automated integration tests for this module are currently non-functional or
-failing.
-{{< /notice >}}
-
 ## About
 
 A `dgraph-dql` tool executes a pre-defined DQL statement against a Dgraph

internal/prebuiltconfigs/tools/cloud-sql-mssql-admin.yaml

@@ -46,9 +46,6 @@ tools:
   create_backup:
     kind: cloud-sql-create-backup
     source: cloud-sql-admin-source
-  restore_backup:
-    kind: cloud-sql-restore-backup
-    source: cloud-sql-admin-source
 
 toolsets:
   cloud_sql_mssql_admin_tools:
@@ -61,4 +58,3 @@ toolsets:
     - wait_for_operation
     - clone_instance
     - create_backup
-    - restore_backup

internal/prebuiltconfigs/tools/cloud-sql-mysql-admin.yaml

@@ -46,9 +46,6 @@ tools:
   create_backup:
     kind: cloud-sql-create-backup
     source: cloud-sql-admin-source
-  restore_backup:
-    kind: cloud-sql-restore-backup
-    source: cloud-sql-admin-source
 
 toolsets:
   cloud_sql_mysql_admin_tools:
@@ -61,4 +58,3 @@ toolsets:
     - wait_for_operation
     - clone_instance
     - create_backup
-    - restore_backup

internal/prebuiltconfigs/tools/cloud-sql-postgres-admin.yaml

@@ -49,9 +49,6 @@ tools:
   create_backup:
     kind: cloud-sql-create-backup
     source: cloud-sql-admin-source
-  restore_backup:
-    kind: cloud-sql-restore-backup
-    source: cloud-sql-admin-source
 
 toolsets:
   cloud_sql_postgres_admin_tools:
@@ -65,4 +62,3 @@ toolsets:
     - postgres_upgrade_precheck
     - clone_instance
     - create_backup
-    - restore_backup

internal/server/mcp/v20241105/method.go

@@ -183,6 +183,13 @@ func toolsCallHandler(ctx context.Context, id jsonrpc.RequestId, resourceMgr *re
 	}
 	logger.DebugContext(ctx, fmt.Sprintf("invocation params: %s", params))
 
+	embeddingModels := resourceMgr.GetEmbeddingModelMap()
+	params, err = tool.EmbedParams(ctx, params, embeddingModels)
+	if err != nil {
+		err = fmt.Errorf("error embedding parameters: %w", err)
+		return jsonrpc.NewError(id, jsonrpc.INVALID_PARAMS, err.Error(), nil), err
+	}
+
 	// run tool invocation and generate response.
 	results, err := tool.Invoke(ctx, resourceMgr, params, accessToken)
 	if err != nil {

internal/server/mcp/v20250326/method.go

@@ -183,6 +183,13 @@ func toolsCallHandler(ctx context.Context, id jsonrpc.RequestId, resourceMgr *re
 	}
 	logger.DebugContext(ctx, fmt.Sprintf("invocation params: %s", params))
 
+	embeddingModels := resourceMgr.GetEmbeddingModelMap()
+	params, err = tool.EmbedParams(ctx, params, embeddingModels)
+	if err != nil {
+		err = fmt.Errorf("error embedding parameters: %w", err)
+		return jsonrpc.NewError(id, jsonrpc.INVALID_PARAMS, err.Error(), nil), err
+	}
+
 	// run tool invocation and generate response.
 	results, err := tool.Invoke(ctx, resourceMgr, params, accessToken)
 	if err != nil {

internal/server/mcp/v20250618/method.go

@@ -176,6 +176,13 @@ func toolsCallHandler(ctx context.Context, id jsonrpc.RequestId, resourceMgr *re
 	}
 	logger.DebugContext(ctx, fmt.Sprintf("invocation params: %s", params))
 
+	embeddingModels := resourceMgr.GetEmbeddingModelMap()
+	params, err = tool.EmbedParams(ctx, params, embeddingModels)
+	if err != nil {
+		err = fmt.Errorf("error embedding parameters: %w", err)
+		return jsonrpc.NewError(id, jsonrpc.INVALID_PARAMS, err.Error(), nil), err
+	}
+
 	// run tool invocation and generate response.
 	results, err := tool.Invoke(ctx, resourceMgr, params, accessToken)
 	if err != nil {

internal/sources/cloudsqladmin/cloud_sql_admin.go

@@ -19,7 +19,6 @@ import (
 	"fmt"
 	"net/http"
 	"regexp"
-	"strconv"
 	"strings"
 	"text/template"
 	"time"
@@ -37,10 +36,7 @@ import (
 
 const SourceKind string = "cloud-sql-admin"
 
-var (
-	targetLinkRegex = regexp.MustCompile(`/projects/([^/]+)/instances/([^/]+)/databases/([^/]+)`)
-	backupDRRegex   = regexp.MustCompile(`^projects/([^/]+)/locations/([^/]+)/backupVaults/([^/]+)/dataSources/([^/]+)/backups/([^/]+)$`)
-)
+var targetLinkRegex = regexp.MustCompile(`/projects/([^/]+)/instances/([^/]+)/databases/([^/]+)`)
 
 // validate interface
 var _ sources.SourceConfig = Config{}
@@ -378,48 +374,6 @@ func (s *Source) InsertBackupRun(ctx context.Context, project, instance, locatio
 	return resp, nil
 }
 
-func (s *Source) RestoreBackup(ctx context.Context, targetProject, targetInstance, sourceProject, sourceInstance, backupID, accessToken string) (any, error) {
-	request := &sqladmin.InstancesRestoreBackupRequest{}
-
-	// There are 3 scenarios for the backup identifier:
-	// 1. The identifier is an int64 containing the timestamp of the BackupRun.
-	//    This is used to restore standard backups, and the RestoreBackupContext
-	//    field should be populated with the backup ID and source instance info.
-	// 2. The identifier is a string of the format
-	//    'projects/{project-id}/locations/{location}/backupVaults/{backupvault}/dataSources/{datasource}/backups/{backup-uid}'.
-	//    This is used to restore BackupDR backups, and the BackupdrBackup field
-	//    should be populated.
-	// 3. The identifer is a string of the format
-	//    'projects/{project-id}/backups/{backup-uid}'. In this case, the Backup
-	//    field should be populated.
-	if backupRunID, err := strconv.ParseInt(backupID, 10, 64); err == nil {
-		if sourceProject == "" || targetInstance == "" {
-			return nil, fmt.Errorf("source project and instance are required when restoring via backup ID")
-		}
-		request.RestoreBackupContext = &sqladmin.RestoreBackupContext{
-			Project:     sourceProject,
-			InstanceId:  sourceInstance,
-			BackupRunId: backupRunID,
-		}
-	} else if backupDRRegex.MatchString(backupID) {
-		request.BackupdrBackup = backupID
-	} else {
-		request.Backup = backupID
-	}
-
-	service, err := s.GetService(ctx, string(accessToken))
-	if err != nil {
-		return nil, err
-	}
-
-	resp, err := service.Instances.RestoreBackup(targetProject, targetInstance, request).Do()
-	if err != nil {
-		return nil, fmt.Errorf("error restoring backup: %w", err)
-	}
-
-	return resp, nil
-}
-
 func generateCloudSQLConnectionMessage(ctx context.Context, source *Source, logger log.Logger, opResponse map[string]any, connectionMessageTemplate string) (string, bool) {
 	operationType, ok := opResponse["operationType"].(string)
 	if !ok || operationType != "CREATE_DATABASE" {

internal/tools/cloudsql/cloudsqlrestorebackup/cloudsqlrestorebackup.go

@@ -1,183 +0,0 @@
-// Copyright 2026 Google LLC
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//      http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-package cloudsqlrestorebackup
-
-import (
-	"context"
-	"fmt"
-
-	"github.com/goccy/go-yaml"
-	"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
-	"github.com/googleapis/genai-toolbox/internal/sources"
-	"github.com/googleapis/genai-toolbox/internal/tools"
-	"github.com/googleapis/genai-toolbox/internal/util/parameters"
-	"google.golang.org/api/sqladmin/v1"
-)
-
-const kind string = "cloud-sql-restore-backup"
-
-var _ tools.ToolConfig = Config{}
-
-type compatibleSource interface {
-	GetDefaultProject() string
-	GetService(context.Context, string) (*sqladmin.Service, error)
-	UseClientAuthorization() bool
-	RestoreBackup(ctx context.Context, targetProject, targetInstance, sourceProject, sourceInstance, backupID, accessToken string) (any, error)
-}
-
-// Config defines the configuration for the restore-backup tool.
-type Config struct {
-	Name         string   `yaml:"name" validate:"required"`
-	Kind         string   `yaml:"kind" validate:"required"`
-	Description  string   `yaml:"description"`
-	Source       string   `yaml:"source" validate:"required"`
-	AuthRequired []string `yaml:"authRequired"`
-}
-
-func init() {
-	if !tools.Register(kind, newConfig) {
-		panic(fmt.Sprintf("tool kind %q already registered", kind))
-	}
-}
-
-func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) {
-	actual := Config{Name: name}
-	if err := decoder.DecodeContext(ctx, &actual); err != nil {
-		return nil, err
-	}
-	return actual, nil
-}
-
-// ToolConfigKind returns the kind of the tool.
-func (cfg Config) ToolConfigKind() string {
-	return kind
-}
-
-// Initialize initializes the tool from the configuration.
-func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) {
-	rawS, ok := srcs[cfg.Source]
-	if !ok {
-		return nil, fmt.Errorf("no source named %q configured", cfg.Source)
-	}
-	s, ok := rawS.(compatibleSource)
-	if !ok {
-		return nil, fmt.Errorf("invalid source for %q tool: source %q not compatible", kind, cfg.Source)
-	}
-
-	project := s.GetDefaultProject()
-	var targetProjectParam parameters.Parameter
-	if project != "" {
-		targetProjectParam = parameters.NewStringParameterWithDefault("target_project", project, "The GCP project ID. This is pre-configured; do not ask for it unless the user explicitly provides a different one.")
-	} else {
-		targetProjectParam = parameters.NewStringParameter("target_project", "The project ID")
-	}
-
-	allParameters := parameters.Parameters{
-		targetProjectParam,
-		parameters.NewStringParameter("target_instance", "Cloud SQL instance ID of the target instance. This does not include the project ID."),
-		parameters.NewStringParameter("backup_id", "Identifier of the backup being restored. Can be a BackupRun ID, backup name, or BackupDR backup name. Use the full backup ID as provided, do not try to parse it"),
-		parameters.NewStringParameterWithRequired("source_project", "GCP project ID of the instance that the backup belongs to. Only required if the backup_id is a BackupRun ID.", false),
-		parameters.NewStringParameterWithRequired("source_instance", "Cloud SQL instance ID of the instance that the backup belongs to. Only required if the backup_id is a BackupRun ID.", false),
-	}
-	paramManifest := allParameters.Manifest()
-
-	description := cfg.Description
-	if description == "" {
-		description = "Restores a backup on a Cloud SQL instance."
-	}
-
-	mcpManifest := tools.GetMcpManifest(cfg.Name, description, cfg.AuthRequired, allParameters, nil)
-
-	return Tool{
-		Config:      cfg,
-		AllParams:   allParameters,
-		manifest:    tools.Manifest{Description: description, Parameters: paramManifest, AuthRequired: cfg.AuthRequired},
-		mcpManifest: mcpManifest,
-	}, nil
-}
-
-// Tool represents the restore-backup tool.
-type Tool struct {
-	Config
-	AllParams   parameters.Parameters `yaml:"allParams"`
-	manifest    tools.Manifest
-	mcpManifest tools.McpManifest
-}
-
-func (t Tool) ToConfig() tools.ToolConfig {
-	return t.Config
-}
-
-func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) {
-	source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Kind)
-	if err != nil {
-		return nil, err
-	}
-	paramsMap := params.AsMap()
-
-	targetProject, ok := paramsMap["target_project"].(string)
-	if !ok {
-		return nil, fmt.Errorf("error casting 'target_project' parameter: %v", paramsMap["target_project"])
-	}
-	targetInstance, ok := paramsMap["target_instance"].(string)
-	if !ok {
-		return nil, fmt.Errorf("error casting 'target_instance' parameter: %v", paramsMap["target_instance"])
-	}
-	backupID, ok := paramsMap["backup_id"].(string)
-	if !ok {
-		return nil, fmt.Errorf("error casting 'backup_id' parameter: %v", paramsMap["backup_id"])
-	}
-	sourceProject, _ := paramsMap["source_project"].(string)
-	sourceInstance, _ := paramsMap["source_instance"].(string)
-
-	return source.RestoreBackup(ctx, targetProject, targetInstance, sourceProject, sourceInstance, backupID, string(accessToken))
-}
-
-// ParseParams parses the parameters for the tool.
-func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (parameters.ParamValues, error) {
-	return parameters.ParseParams(t.AllParams, data, claims)
-}
-
-func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {
-	return parameters.EmbedParams(ctx, t.AllParams, paramValues, embeddingModelsMap, nil)
-}
-
-// Manifest returns the tool's manifest.
-func (t Tool) Manifest() tools.Manifest {
-	return t.manifest
-}
-
-// McpManifest returns the tool's MCP manifest.
-func (t Tool) McpManifest() tools.McpManifest {
-	return t.mcpManifest
-}
-
-// Authorized checks if the tool is authorized.
-func (t Tool) Authorized(verifiedAuthServices []string) bool {
-	return true
-}
-
-func (t Tool) RequiresClientAuthorization(resourceMgr tools.SourceProvider) (bool, error) {
-	source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Kind)
-	if err != nil {
-		return false, err
-	}
-
-	return source.UseClientAuthorization(), nil
-}
-
-func (t Tool) GetAuthTokenHeaderName(resourceMgr tools.SourceProvider) (string, error) {
-	return "Authorization", nil
-}

internal/tools/cloudsql/cloudsqlrestorebackup/cloudsqlrestorebackup_test.go

@@ -1,71 +0,0 @@
-// Copyright 2026 Google LLC
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//      http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-package cloudsqlrestorebackup_test
-
-import (
-	"testing"
-
-	yaml "github.com/goccy/go-yaml"
-	"github.com/google/go-cmp/cmp"
-	"github.com/googleapis/genai-toolbox/internal/server"
-	"github.com/googleapis/genai-toolbox/internal/testutils"
-	"github.com/googleapis/genai-toolbox/internal/tools/cloudsql/cloudsqlrestorebackup"
-)
-
-func TestParseFromYaml(t *testing.T) {
-	ctx, err := testutils.ContextWithNewLogger()
-	if err != nil {
-		t.Fatalf("unexpected error: %s", err)
-	}
-	tcs := []struct {
-		desc string
-		in   string
-		want server.ToolConfigs
-	}{
-		{
-			desc: "basic example",
-			in: `
-			tools:
-			  restore-backup-tool:
-			    kind: cloud-sql-restore-backup
-			    description: a test description
-			    source: a-source
-			`,
-			want: server.ToolConfigs{
-				"restore-backup-tool": cloudsqlrestorebackup.Config{
-					Name:         "restore-backup-tool",
-					Kind:         "cloud-sql-restore-backup",
-					Description:  "a test description",
-					Source:       "a-source",
-					AuthRequired: []string{},
-				},
-			},
-		},
-	}
-	for _, tc := range tcs {
-		t.Run(tc.desc, func(t *testing.T) {
-			got := struct {
-				Tools server.ToolConfigs `yaml:"tools"`
-			}{}
-			err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got)
-			if err != nil {
-				t.Fatalf("unable to unmarshal: %s", err)
-			}
-			if diff := cmp.Diff(tc.want, got.Tools); diff != "" {
-				t.Fatalf("incorrect parse: diff %v", diff)
-			}
-		})
-	}
-}

tests/alloydbpg/alloydb_pg_integration_test.go

@@ -147,12 +147,20 @@ func TestAlloyDBPgToolEndpoints(t *testing.T) {
 	teardownTable2 := tests.SetupPostgresSQLTable(t, ctx, pool, createAuthTableStmt, insertAuthTableStmt, tableNameAuth, authTestParams)
 	defer teardownTable2(t)
 
+	// Set up table for semanti search
+	vectorTableName, tearDownVectorTable := tests.SetupPostgresVectorTable(t, ctx, pool)
+	defer tearDownVectorTable(t)
+
 	// Write config into a file and pass it to command
 	toolsFile := tests.GetToolsConfig(sourceConfig, AlloyDBPostgresToolKind, paramToolStmt, idParamToolStmt, nameParamToolStmt, arrayToolStmt, authToolStmt)
 	toolsFile = tests.AddExecuteSqlConfig(t, toolsFile, "postgres-execute-sql")
 	tmplSelectCombined, tmplSelectFilterCombined := tests.GetPostgresSQLTmplToolStatement()
 	toolsFile = tests.AddTemplateParamConfig(t, toolsFile, AlloyDBPostgresToolKind, tmplSelectCombined, tmplSelectFilterCombined, "")
 
+	// Add semantic search tool config
+	insertStmt, searchStmt := tests.GetPostgresVectorSearchStmts(vectorTableName)
+	toolsFile = tests.AddSemanticSearchConfig(t, toolsFile, AlloyDBPostgresToolKind, insertStmt, searchStmt)
+
 	toolsFile = tests.AddPostgresPrebuiltConfig(t, toolsFile)
 
 	cmd, cleanup, err := tests.StartCmd(ctx, toolsFile, args...)

tests/cloudsql/cloud_sql_restore_backup_test.go

@@ -1,267 +0,0 @@
-// Copyright 2026 Google LLC
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//      http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-package cloudsql
-
-import (
-	"bytes"
-	"context"
-	"encoding/json"
-	"fmt"
-	"io"
-	"net/http"
-	"net/http/httptest"
-	"net/url"
-	"reflect"
-	"regexp"
-	"strings"
-	"testing"
-	"time"
-
-	"github.com/google/go-cmp/cmp"
-	"github.com/googleapis/genai-toolbox/internal/testutils"
-	"github.com/googleapis/genai-toolbox/tests"
-	"google.golang.org/api/sqladmin/v1"
-)
-
-var (
-	restoreBackupToolKind = "cloud-sql-restore-backup"
-)
-
-type restoreBackupTransport struct {
-	transport http.RoundTripper
-	url       *url.URL
-}
-
-func (t *restoreBackupTransport) RoundTrip(req *http.Request) (*http.Response, error) {
-	if strings.HasPrefix(req.URL.String(), "https://sqladmin.googleapis.com") {
-		req.URL.Scheme = t.url.Scheme
-		req.URL.Host = t.url.Host
-	}
-	return t.transport.RoundTrip(req)
-}
-
-type masterRestoreBackupHandler struct {
-	t *testing.T
-}
-
-func (h *masterRestoreBackupHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
-	if !strings.Contains(r.UserAgent(), "genai-toolbox/") {
-		h.t.Errorf("User-Agent header not found")
-	}
-	var body sqladmin.InstancesRestoreBackupRequest
-	if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
-		h.t.Fatalf("failed to decode request body: %v", err)
-	} else {
-		h.t.Logf("Received request body: %+v", body)
-	}
-
-	var expectedBody sqladmin.InstancesRestoreBackupRequest
-	var response any
-	var statusCode int
-
-	switch {
-	case body.Backup != "":
-		expectedBody = sqladmin.InstancesRestoreBackupRequest{
-			Backup: "projects/p1/backups/test-uid",
-		}
-		response = map[string]any{"name": "op1", "status": "PENDING"}
-		statusCode = http.StatusOK
-	case body.BackupdrBackup != "":
-		expectedBody = sqladmin.InstancesRestoreBackupRequest{
-			BackupdrBackup: "projects/p1/locations/us-central1/backupVaults/test-vault/dataSources/test-ds/backups/test-uid",
-		}
-		response = map[string]any{"name": "op1", "status": "PENDING"}
-		statusCode = http.StatusOK
-	case body.RestoreBackupContext != nil:
-		expectedBody = sqladmin.InstancesRestoreBackupRequest{
-			RestoreBackupContext: &sqladmin.RestoreBackupContext{
-				Project:     "p1",
-				InstanceId:  "source",
-				BackupRunId: 12345,
-			},
-		}
-		response = map[string]any{"name": "op1", "status": "PENDING"}
-		statusCode = http.StatusOK
-	default:
-		http.Error(w, fmt.Sprintf("unhandled restore request body: %v", body), http.StatusInternalServerError)
-		return
-	}
-
-	if diff := cmp.Diff(expectedBody, body); diff != "" {
-		h.t.Errorf("unexpected request body (-want +got):\n%s", diff)
-	}
-
-	w.Header().Set("Content-Type", "application/json")
-	w.WriteHeader(statusCode)
-	if err := json.NewEncoder(w).Encode(response); err != nil {
-		http.Error(w, err.Error(), http.StatusInternalServerError)
-	}
-}
-
-func TestRestoreBackupToolEndpoints(t *testing.T) {
-	ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
-	defer cancel()
-
-	handler := &masterRestoreBackupHandler{t: t}
-	server := httptest.NewServer(handler)
-	defer server.Close()
-
-	serverURL, err := url.Parse(server.URL)
-	if err != nil {
-		t.Fatalf("failed to parse server URL: %v", err)
-	}
-
-	originalTransport := http.DefaultClient.Transport
-	if originalTransport == nil {
-		originalTransport = http.DefaultTransport
-	}
-	http.DefaultClient.Transport = &restoreBackupTransport{
-		transport: originalTransport,
-		url:       serverURL,
-	}
-	t.Cleanup(func() {
-		http.DefaultClient.Transport = originalTransport
-	})
-
-	var args []string
-	toolsFile := getRestoreBackupToolsConfig()
-	cmd, cleanup, err := tests.StartCmd(ctx, toolsFile, args...)
-	if err != nil {
-		t.Fatalf("command initialization returned an error: %s", err)
-	}
-	defer cleanup()
-
-	waitCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
-	defer cancel()
-	out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
-	if err != nil {
-		t.Logf("toolbox command logs: \n%s", out)
-		t.Fatalf("toolbox didn't start successfully: %s", err)
-	}
-
-	tcs := []struct {
-		name        string
-		toolName    string
-		body        string
-		want        string
-		expectError bool
-		errorStatus int
-	}{
-		{
-			name:     "successful restore with standard backup",
-			toolName: "restore-backup",
-			body:     `{"target_project": "p1", "target_instance": "instance-standard", "backup_id": "12345", "source_project": "p1", "source_instance": "source"}`,
-			want:     `{"name":"op1","status":"PENDING"}`,
-		},
-		{
-			name:     "successful restore with project level backup",
-			toolName: "restore-backup",
-			body:     `{"target_project": "p1", "target_instance": "instance-project-level", "backup_id": "projects/p1/backups/test-uid"}`,
-			want:     `{"name":"op1","status":"PENDING"}`,
-		},
-		{
-			name:     "successful restore with BackupDR backup",
-			toolName: "restore-backup",
-			body:     `{"target_project": "p1", "target_instance": "instance-project-level", "backup_id": "projects/p1/locations/us-central1/backupVaults/test-vault/dataSources/test-ds/backups/test-uid"}`,
-			want:     `{"name":"op1","status":"PENDING"}`,
-		},
-		{
-			name:        "missing source instance info for standard backup",
-			toolName:    "restore-backup",
-			body:        `{"target_project": "p1", "target_instance": "instance-project-level", "backup_id": "12345"}`,
-			expectError: true,
-			errorStatus: http.StatusBadRequest,
-		},
-		{
-			name:        "missing backup identifier",
-			toolName:    "restore-backup",
-			body:        `{"target_project": "p1", "target_instance": "instance-project-level"}`,
-			expectError: true,
-			errorStatus: http.StatusBadRequest,
-		},
-		{
-			name:        "missing target instance info",
-			toolName:    "restore-backup",
-			body:        `{"backup_id": "12345"}`,
-			expectError: true,
-			errorStatus: http.StatusBadRequest,
-		},
-	}
-
-	for _, tc := range tcs {
-		tc := tc
-		t.Run(tc.name, func(t *testing.T) {
-			api := fmt.Sprintf("http://127.0.0.1:5000/api/tool/%s/invoke", tc.toolName)
-			req, err := http.NewRequest(http.MethodPost, api, bytes.NewBufferString(tc.body))
-			if err != nil {
-				t.Fatalf("unable to create request: %s", err)
-			}
-			req.Header.Add("Content-type", "application/json")
-			resp, err := http.DefaultClient.Do(req)
-			if err != nil {
-				t.Fatalf("unable to send request: %s", err)
-			}
-			defer resp.Body.Close()
-
-			if tc.expectError {
-				if resp.StatusCode != tc.errorStatus {
-					bodyBytes, _ := io.ReadAll(resp.Body)
-					t.Fatalf("expected status %d but got %d: %s", tc.errorStatus, resp.StatusCode, string(bodyBytes))
-				}
-				return
-			}
-
-			if resp.StatusCode != http.StatusOK {
-				bodyBytes, _ := io.ReadAll(resp.Body)
-				t.Fatalf("response status code is not 200, got %d: %s", resp.StatusCode, string(bodyBytes))
-			}
-
-			var result struct {
-				Result string `json:"result"`
-			}
-			if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
-				t.Fatalf("failed to decode response: %v", err)
-			}
-
-			var got, want map[string]any
-			if err := json.Unmarshal([]byte(result.Result), &got); err != nil {
-				t.Fatalf("failed to unmarshal result: %v", err)
-			}
-			if err := json.Unmarshal([]byte(tc.want), &want); err != nil {
-				t.Fatalf("failed to unmarshal want: %v", err)
-			}
-
-			if !reflect.DeepEqual(got, want) {
-				t.Fatalf("unexpected result: got %+v, want %+v", got, want)
-			}
-		})
-	}
-}
-
-func getRestoreBackupToolsConfig() map[string]any {
-	return map[string]any{
-		"sources": map[string]any{
-			"my-cloud-sql-source": map[string]any{
-				"kind": "cloud-sql-admin",
-			},
-		},
-		"tools": map[string]any{
-			"restore-backup": map[string]any{
-				"kind":   restoreBackupToolKind,
-				"source": "my-cloud-sql-source",
-			},
-		},
-	}
-}

tests/cloudsqlpg/cloud_sql_pg_integration_test.go

@@ -132,12 +132,20 @@ func TestCloudSQLPgSimpleToolEndpoints(t *testing.T) {
 	teardownTable2 := tests.SetupPostgresSQLTable(t, ctx, pool, createAuthTableStmt, insertAuthTableStmt, tableNameAuth, authTestParams)
 	defer teardownTable2(t)
 
+	// Set up table for semantic search
+	vectorTableName, tearDownVectorTable := tests.SetupPostgresVectorTable(t, ctx, pool)
+	defer tearDownVectorTable(t)
+
 	// Write config into a file and pass it to command
 	toolsFile := tests.GetToolsConfig(sourceConfig, CloudSQLPostgresToolKind, paramToolStmt, idParamToolStmt, nameParamToolStmt, arrayToolStmt, authToolStmt)
 	toolsFile = tests.AddExecuteSqlConfig(t, toolsFile, "postgres-execute-sql")
 	tmplSelectCombined, tmplSelectFilterCombined := tests.GetPostgresSQLTmplToolStatement()
 	toolsFile = tests.AddTemplateParamConfig(t, toolsFile, CloudSQLPostgresToolKind, tmplSelectCombined, tmplSelectFilterCombined, "")
 
+	// Add semantic search tool config
+	insertStmt, searchStmt := tests.GetPostgresVectorSearchStmts(vectorTableName)
+	toolsFile = tests.AddSemanticSearchConfig(t, toolsFile, CloudSQLPostgresToolKind, insertStmt, searchStmt)
+
 	toolsFile = tests.AddPostgresPrebuiltConfig(t, toolsFile)
 	cmd, cleanup, err := tests.StartCmd(ctx, toolsFile, args...)
 	if err != nil {
@@ -186,6 +194,7 @@ func TestCloudSQLPgSimpleToolEndpoints(t *testing.T) {
 	tests.RunPostgresListDatabaseStatsTest(t, ctx, pool)
 	tests.RunPostgresListRolesTest(t, ctx, pool)
 	tests.RunPostgresListStoredProcedureTest(t, ctx, pool)
+	tests.RunSemanticSearchToolInvokeTest(t, "null", "", "The quick brown fox")
 }
 
 // Test connection with different IP type

tests/embedding.go

@@ -0,0 +1,251 @@
+// Copyright 2026 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Package tests contains end to end tests meant to verify the Toolbox Server
+// works as expected when executed as a binary.
+
+package tests
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"os"
+	"strings"
+	"testing"
+
+	"github.com/google/uuid"
+	"github.com/googleapis/genai-toolbox/internal/server/mcp/jsonrpc"
+	"github.com/jackc/pgx/v5/pgxpool"
+)
+
+var apiKey = os.Getenv("API_KEY")
+
+// AddSemanticSearchConfig adds embedding models and semantic search tools to the config
+// with configurable tool kind and SQL statements.
+func AddSemanticSearchConfig(t *testing.T, config map[string]any, toolKind, insertStmt, searchStmt string) map[string]any {
+	config["embeddingModels"] = map[string]any{
+		"gemini_model": map[string]any{
+			"kind":      "gemini",
+			"model":     "gemini-embedding-001",
+			"apiKey":    apiKey,
+			"dimension": 768,
+		},
+	}
+
+	tools, ok := config["tools"].(map[string]any)
+	if !ok {
+		t.Fatalf("unable to get tools from config")
+	}
+
+	tools["insert_docs"] = map[string]any{
+		"kind":        toolKind,
+		"source":      "my-instance",
+		"description": "Stores content and its vector embedding into the documents table.",
+		"statement":   insertStmt,
+		"parameters": []any{
+			map[string]any{
+				"name":        "content",
+				"type":        "string",
+				"description": "The text content associated with the vector.",
+			},
+			map[string]any{
+				"name":        "text_to_embed",
+				"type":        "string",
+				"description": "The text content used to generate the vector.",
+				"embeddedBy":  "gemini_model",
+			},
+		},
+	}
+
+	tools["search_docs"] = map[string]any{
+		"kind":        toolKind,
+		"source":      "my-instance",
+		"description": "Finds the most semantically similar document to the query vector.",
+		"statement":   searchStmt,
+		"parameters": []any{
+			map[string]any{
+				"name":        "query",
+				"type":        "string",
+				"description": "The text content to search for.",
+				"embeddedBy":  "gemini_model",
+			},
+		},
+	}
+
+	config["tools"] = tools
+	return config
+}
+
+// RunSemanticSearchToolInvokeTest runs the insert_docs and search_docs tools
+// via both HTTP and MCP endpoints and verifies the output.
+func RunSemanticSearchToolInvokeTest(t *testing.T, insertWant, mcpInsertWant, searchWant string) {
+	// Initialize MCP session once for the MCP test cases
+	sessionId := RunInitialize(t, "2024-11-05")
+
+	tcs := []struct {
+		name        string
+		api         string
+		isMcp       bool
+		requestBody interface{}
+		want        string
+	}{
+		{
+			name:        "HTTP invoke insert_docs",
+			api:         "http://127.0.0.1:5000/api/tool/insert_docs/invoke",
+			isMcp:       false,
+			requestBody: `{"content": "The quick brown fox jumps over the lazy dog", "text_to_embed": "The quick brown fox jumps over the lazy dog"}`,
+			want:        insertWant,
+		},
+		{
+			name:        "HTTP invoke search_docs",
+			api:         "http://127.0.0.1:5000/api/tool/search_docs/invoke",
+			isMcp:       false,
+			requestBody: `{"query": "fast fox jumping"}`,
+			want:        searchWant,
+		},
+		{
+			name:  "MCP invoke insert_docs",
+			api:   "http://127.0.0.1:5000/mcp",
+			isMcp: true,
+			requestBody: jsonrpc.JSONRPCRequest{
+				Jsonrpc: "2.0",
+				Id:      "mcp-insert-docs",
+				Request: jsonrpc.Request{
+					Method: "tools/call",
+				},
+				Params: map[string]any{
+					"name": "insert_docs",
+					"arguments": map[string]any{
+						"content":       "The quick brown fox jumps over the lazy dog",
+						"text_to_embed": "The quick brown fox jumps over the lazy dog",
+					},
+				},
+			},
+			want: mcpInsertWant,
+		},
+		{
+			name:  "MCP invoke search_docs",
+			api:   "http://127.0.0.1:5000/mcp",
+			isMcp: true,
+			requestBody: jsonrpc.JSONRPCRequest{
+				Jsonrpc: "2.0",
+				Id:      "mcp-search-docs",
+				Request: jsonrpc.Request{
+					Method: "tools/call",
+				},
+				Params: map[string]any{
+					"name": "search_docs",
+					"arguments": map[string]any{
+						"query": "fast fox jumping",
+					},
+				},
+			},
+			want: searchWant,
+		},
+	}
+
+	for _, tc := range tcs {
+		t.Run(tc.name, func(t *testing.T) {
+			var bodyReader io.Reader
+			headers := map[string]string{}
+
+			// Prepare Request Body and Headers
+			if tc.isMcp {
+				reqBytes, err := json.Marshal(tc.requestBody)
+				if err != nil {
+					t.Fatalf("failed to marshal mcp request: %v", err)
+				}
+				bodyReader = bytes.NewBuffer(reqBytes)
+				if sessionId != "" {
+					headers["Mcp-Session-Id"] = sessionId
+				}
+			} else {
+				bodyReader = bytes.NewBufferString(tc.requestBody.(string))
+			}
+
+			// Send Request
+			resp, respBody := RunRequest(t, http.MethodPost, tc.api, bodyReader, headers)
+
+			if resp.StatusCode != http.StatusOK {
+				t.Fatalf("response status code is not 200, got %d: %s", resp.StatusCode, string(respBody))
+			}
+
+			// Normalize Response to get the actual tool result string
+			var got string
+			if tc.isMcp {
+				var mcpResp struct {
+					Result struct {
+						Content []struct {
+							Text string `json:"text"`
+						} `json:"content"`
+					} `json:"result"`
+				}
+				if err := json.Unmarshal(respBody, &mcpResp); err != nil {
+					t.Fatalf("error parsing mcp response: %s", err)
+				}
+				if len(mcpResp.Result.Content) > 0 {
+					got = mcpResp.Result.Content[0].Text
+				}
+			} else {
+				var httpResp map[string]interface{}
+				if err := json.Unmarshal(respBody, &httpResp); err != nil {
+					t.Fatalf("error parsing http response: %s", err)
+				}
+				if res, ok := httpResp["result"].(string); ok {
+					got = res
+				}
+			}
+
+			if !strings.Contains(got, tc.want) {
+				t.Fatalf("unexpected value: got %q, want %q", got, tc.want)
+			}
+		})
+	}
+}
+
+// SetupPostgresVectorTable sets up the vector extension and a vector table
+func SetupPostgresVectorTable(t *testing.T, ctx context.Context, pool *pgxpool.Pool) (string, func(*testing.T)) {
+	t.Helper()
+	if _, err := pool.Exec(ctx, "CREATE EXTENSION IF NOT EXISTS vector"); err != nil {
+		t.Fatalf("failed to create vector extension: %v", err)
+	}
+
+	tableName := "vector_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
+
+	createTableStmt := fmt.Sprintf(`CREATE TABLE %s (
+		id SERIAL PRIMARY KEY, 
+		content TEXT, 
+		embedding vector(768)
+	)`, tableName)
+
+	if _, err := pool.Exec(ctx, createTableStmt); err != nil {
+		t.Fatalf("failed to create table %s: %v", tableName, err)
+	}
+
+	return tableName, func(t *testing.T) {
+		if _, err := pool.Exec(ctx, fmt.Sprintf("DROP TABLE IF EXISTS %s", tableName)); err != nil {
+			t.Errorf("failed to drop table %s: %v", tableName, err)
+		}
+	}
+}
+
+func GetPostgresVectorSearchStmts(vectorTableName string) (string, string) {
+	insertStmt := fmt.Sprintf("INSERT INTO %s (content, embedding) VALUES ($1, $2)", vectorTableName)
+	searchStmt := fmt.Sprintf("SELECT id, content, embedding <-> $1 AS distance FROM %s ORDER BY distance LIMIT 1", vectorTableName)
+	return insertStmt, searchStmt
+}

tests/postgres/postgres_integration_test.go

@@ -111,6 +111,10 @@ func TestPostgres(t *testing.T) {
 	teardownTable2 := tests.SetupPostgresSQLTable(t, ctx, pool, createAuthTableStmt, insertAuthTableStmt, tableNameAuth, authTestParams)
 	defer teardownTable2(t)
 
+	// Set up table for semantic search
+	vectorTableName, tearDownVectorTable := tests.SetupPostgresVectorTable(t, ctx, pool)
+	defer tearDownVectorTable(t)
+
 	// Write config into a file and pass it to command
 	toolsFile := tests.GetToolsConfig(sourceConfig, PostgresToolKind, paramToolStmt, idParamToolStmt, nameParamToolStmt, arrayToolStmt, authToolStmt)
 	toolsFile = tests.AddExecuteSqlConfig(t, toolsFile, "postgres-execute-sql")
@@ -118,6 +122,10 @@ func TestPostgres(t *testing.T) {
 	toolsFile = tests.AddTemplateParamConfig(t, toolsFile, PostgresToolKind, tmplSelectCombined, tmplSelectFilterCombined, "")
 	toolsFile = tests.AddPostgresPrebuiltConfig(t, toolsFile)
 
+	// Add semantic search tool config
+	insertStmt, searchStmt := tests.GetPostgresVectorSearchStmts(vectorTableName)
+	toolsFile = tests.AddSemanticSearchConfig(t, toolsFile, PostgresToolKind, insertStmt, searchStmt)
+
 	cmd, cleanup, err := tests.StartCmd(ctx, toolsFile, args...)
 	if err != nil {
 		t.Fatalf("command initialization returned an error: %s", err)
@@ -165,4 +173,5 @@ func TestPostgres(t *testing.T) {
 	tests.RunPostgresListDatabaseStatsTest(t, ctx, pool)
 	tests.RunPostgresListRolesTest(t, ctx, pool)
 	tests.RunPostgresListStoredProcedureTest(t, ctx, pool)
+	tests.RunSemanticSearchToolInvokeTest(t, "null", "", "The quick brown fox")
 }