Compare commits

..

7 Commits

Author SHA1 Message Date
Juexin Wang
e716efea6e Refactor: move getCloudGdaToolsConfig to top of file for readability 2026-01-16 17:26:14 -08:00
Juexin Wang
dcfd056a30 Fix lint errors: format tests/cloudgda/cloud_gda_integration_test.go with goimports 2026-01-16 17:18:31 -08:00
Juexin Wang
59aa21729d Update go.mod: promote grpc to direct dependency for integration tests 2026-01-16 17:10:06 -08:00
Juexin Wang
e25ee6f165 Refactor Cloud GDA source to use package-level variable for testing and use gRPC mock server in integration tests 2026-01-16 17:08:41 -08:00
Juexin Wang
9bb4eee494 refactor(tools/cloudgda): update to use google-cloud-go sdk types
Removes types.go and uses geminidataanalyticspb types with wrapper structs for YAML decoding.
2026-01-15 10:24:10 -08:00
Juexin Wang
9f5b04cf73 Refactor Cloud GDA source to support per-request client authorization 2026-01-14 16:10:50 -08:00
Juexin Wang
66d6b58c4f intorduce Go SDK for GDA tool and source 2026-01-14 16:10:50 -08:00
62 changed files with 385 additions and 2470 deletions

View File

@@ -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", "API_KEY"]
["CLOUD_SQL_POSTGRES_USER", "CLOUD_SQL_POSTGRES_PASS", "CLIENT_ID"]
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", "API_KEY"]
secretEnv: ["ALLOYDB_POSTGRES_USER", "ALLOYDB_POSTGRES_PASS", "CLIENT_ID"]
volumes:
- name: "go"
path: "/gopath"
@@ -293,7 +293,7 @@ steps:
.ci/test_with_coverage.sh \
"Cloud Healthcare API" \
cloudhealthcare \
cloudhealthcare
cloudhealthcare || echo "Integration tests failed."
- id: "postgres"
name: golang:1
@@ -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", "API_KEY"]
secretEnv: ["POSTGRES_USER", "POSTGRES_PASS", "CLIENT_ID"]
volumes:
- name: "go"
path: "/gopath"
@@ -964,13 +964,6 @@ 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
@@ -987,6 +980,8 @@ 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

View File

@@ -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"
@@ -386,7 +385,6 @@ 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.UserAgentMetadata, "user-agent-metadata", []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) }

View File

@@ -70,9 +70,6 @@ func withDefaults(c server.ServerConfig) server.ServerConfig {
if c.AllowedHosts == nil {
c.AllowedHosts = []string{"*"}
}
if c.UserAgentMetadata == nil {
c.UserAgentMetadata = []string{}
}
return c
}
@@ -233,13 +230,6 @@ func TestServerConfigFlags(t *testing.T) {
AllowedHosts: []string{"http://foo.com", "http://bar.com"},
}),
},
{
desc: "user agent metadata",
args: []string{"--user-agent-metadata", "foo,bar"},
want: withDefaults(server.ServerConfig{
UserAgentMetadata: []string{"foo", "bar"},
}),
},
}
for _, tc := range tcs {
t.Run(tc.desc, func(t *testing.T) {
@@ -1503,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"},
},
},
},
@@ -1513,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"},
},
},
},
@@ -1523,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"},
},
},
},
@@ -2181,115 +2171,3 @@ func TestDefaultToolsFileBehavior(t *testing.T) {
})
}
}
func TestParameterReferenceValidation(t *testing.T) {
ctx, err := testutils.ContextWithNewLogger()
if err != nil {
t.Fatalf("unexpected error: %s", err)
}
// Base template
baseYaml := `
sources:
dummy-source:
kind: http
baseUrl: http://example.com
tools:
test-tool:
kind: postgres-sql
source: dummy-source
description: test tool
statement: SELECT 1;
parameters:
%s`
tcs := []struct {
desc string
params string
wantErr bool
errSubstr string
}{
{
desc: "valid backward reference",
params: `
- name: source_param
type: string
description: source
- name: copy_param
type: string
description: copy
valueFromParam: source_param`,
wantErr: false,
},
{
desc: "valid forward reference (out of order)",
params: `
- name: copy_param
type: string
description: copy
valueFromParam: source_param
- name: source_param
type: string
description: source`,
wantErr: false,
},
{
desc: "invalid missing reference",
params: `
- name: copy_param
type: string
description: copy
valueFromParam: non_existent_param`,
wantErr: true,
errSubstr: "references '\"non_existent_param\"' in the 'valueFromParam' field",
},
{
desc: "invalid self reference",
params: `
- name: myself
type: string
description: self
valueFromParam: myself`,
wantErr: true,
errSubstr: "parameter \"myself\" cannot copy value from itself",
},
{
desc: "multiple valid references",
params: `
- name: a
type: string
description: a
- name: b
type: string
description: b
valueFromParam: a
- name: c
type: string
description: c
valueFromParam: a`,
wantErr: false,
},
}
for _, tc := range tcs {
t.Run(tc.desc, func(t *testing.T) {
// Indent parameters to match YAML structure
yamlContent := fmt.Sprintf(baseYaml, tc.params)
_, err := parseToolsFile(ctx, []byte(yamlContent))
if tc.wantErr {
if err == nil {
t.Fatal("expected error, got nil")
}
if !strings.Contains(err.Error(), tc.errSubstr) {
t.Errorf("error %q does not contain expected substring %q", err.Error(), tc.errSubstr)
}
} else {
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
}
})
}
}

View File

@@ -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

View File

@@ -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

View File

@@ -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

View File

@@ -20,7 +20,6 @@ The native SDKs can be combined with MCP clients in many cases.
Toolbox currently supports the following versions of MCP specification:
* [2025-11-25](https://modelcontextprotocol.io/specification/2025-11-25)
* [2025-06-18](https://modelcontextprotocol.io/specification/2025-06-18)
* [2025-03-26](https://modelcontextprotocol.io/specification/2025-03-26)
* [2024-11-05](https://modelcontextprotocol.io/specification/2024-11-05)

View File

@@ -207,7 +207,6 @@ You can connect to Toolbox Cloud Run instances directly through the SDK.
{{< tab header="Python" lang="python" >}}
import asyncio
from toolbox_core import ToolboxClient, auth_methods
from toolbox_core.protocol import Protocol
# Replace with the Cloud Run service URL generated in the previous step
URL = "https://cloud-run-url.app"
@@ -218,7 +217,6 @@ async def main():
async with ToolboxClient(
URL,
client_headers={"Authorization": auth_token_provider},
protocol=Protocol.TOOLBOX,
) as toolbox:
toolset = await toolbox.load_toolset()
# ...
@@ -283,5 +281,3 @@ contain the specific error message needed to diagnose the problem.
Manager, it means the Toolbox service account is missing permissions.
- Ensure the `toolbox-identity` service account has the **Secret Manager
Secret Accessor** (`roles/secretmanager.secretAccessor`) IAM role.
- **Cloud Run Connections via IAP:** Currently we do not support Cloud Run connections via [IAP](https://docs.cloud.google.com/iap/docs/concepts-overview). Please disable IAP if you are using it.

View File

@@ -27,7 +27,6 @@ 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

View File

@@ -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

View File

@@ -3,14 +3,13 @@ title: "EmbeddingModels"
type: docs
weight: 2
description: >
EmbeddingModels represent services that transform text into vector embeddings
for semantic search.
EmbeddingModels represent services that transform text into vector embeddings for semantic search.
---
EmbeddingModels represent services that generate vector representations of text
data. In the MCP Toolbox, these models enable **Semantic Queries**, allowing
[Tools](../tools/) to automatically convert human-readable text into numerical
vectors before using them in a query.
data. In the MCP Toolbox, these models enable **Semantic Queries**,
allowing [Tools](../tools/) to automatically convert human-readable text into
numerical vectors before using them in a query.
This is primarily used in two scenarios:
@@ -20,33 +19,14 @@ This is primarily used in two scenarios:
- **Semantic Search**: Converting a natural language query into a vector to
perform similarity searches.
## Hidden Parameter Duplication (valueFromParam)
When building tools for vector ingestion, you often need the same input string
twice:
1. To store the original text in a TEXT column.
1. To generate the vector embedding for a VECTOR column.
Requesting an Agent (LLM) to output the exact same string twice is inefficient
and error-prone. The `valueFromParam` field solves this by allowing a parameter
to inherit its value from another parameter in the same tool.
### Key Behaviors
1. Hidden from Manifest: Parameters with valueFromParam set are excluded from
the tool definition sent to the Agent. The Agent does not know this parameter
exists.
1. Auto-Filled: When the tool is executed, the Toolbox automatically copies the
value from the referenced parameter before processing embeddings.
## Example
The following configuration defines an embedding model and applies it to
specific tool parameters.
{{< notice tip >}} Use environment variable replacement with the format
${ENV_NAME} instead of hardcoding your API keys into the configuration file.
{{< notice tip >}}
Use environment variable replacement with the format ${ENV_NAME}
instead of hardcoding your API keys into the configuration file.
{{< /notice >}}
### Step 1 - Define an Embedding Model
@@ -60,12 +40,14 @@ embeddingModels:
model: gemini-embedding-001
apiKey: ${GOOGLE_API_KEY}
dimension: 768
```
### Step 2 - Embed Tool Parameters
Use the defined embedding model, embed your query parameters using the
`embeddedBy` field. Only string-typed parameters can be embedded:
`embeddedBy` field. Only string-typed
parameters can be embedded:
```yaml
tools:
@@ -79,13 +61,10 @@ tools:
parameters:
- name: content
type: string
description: The raw text content to be stored in the database.
- name: vector_string
type: string
# This parameter is hidden from the LLM.
# It automatically copies the value from 'content' and embeds it.
valueFromParam: content
embeddedBy: gemini-model
description: The text to be vectorized and stored.
embeddedBy: gemini-model # refers to the name of a defined embedding model
# Semantic search tool
search_embedding:

View File

@@ -12,9 +12,6 @@ aliases:
The `cloud-gemini-data-analytics-query` tool allows you to send natural language questions to the Gemini Data Analytics API and receive structured responses containing SQL queries, natural language answers, and explanations. For details on defining data agent context for database data sources, see the official [documentation](https://docs.cloud.google.com/gemini/docs/conversational-analytics-api/data-agent-authored-context-databases).
> [!NOTE]
> Only `alloydb`, `spannerReference`, and `cloudSqlReference` are supported as [datasource references](https://clouddocs.devsite.corp.google.com/gemini/docs/conversational-analytics-api/reference/rest/v1beta/projects.locations.dataAgents#DatasourceReferences).
## Example
```yaml

View File

@@ -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)

View File

@@ -30,10 +30,6 @@ following config for example:
- name: userNames
type: array
description: The user names to be set.
items:
name: userName # the item name doesn't matter but it has to exist
type: string
description: username
```
If the input is an array of strings `["Alice", "Sid", "Bob"]`, The final command

4
go.mod
View File

@@ -12,7 +12,7 @@ require (
cloud.google.com/go/dataplex v1.28.0
cloud.google.com/go/dataproc/v2 v2.15.0
cloud.google.com/go/firestore v1.20.0
cloud.google.com/go/geminidataanalytics v0.3.0
cloud.google.com/go/geminidataanalytics v0.5.0
cloud.google.com/go/longrunning v0.7.0
cloud.google.com/go/spanner v1.86.1
github.com/ClickHouse/clickhouse-go/v2 v2.40.3
@@ -63,6 +63,7 @@ require (
google.golang.org/api v0.256.0
google.golang.org/genai v1.37.0
google.golang.org/genproto v0.0.0-20251022142026-3a174f9686a8
google.golang.org/grpc v1.76.0
google.golang.org/protobuf v1.36.10
modernc.org/sqlite v1.40.0
)
@@ -229,7 +230,6 @@ require (
golang.org/x/xerrors v0.0.0-20240903120638-7835f813f4da // indirect
google.golang.org/genproto/googleapis/api v0.0.0-20251111163417-95abcf5c77ba // indirect
google.golang.org/genproto/googleapis/rpc v0.0.0-20251103181224-f26f9409b101 // indirect
google.golang.org/grpc v1.76.0 // indirect
gopkg.in/inf.v0 v0.9.1 // indirect
gopkg.in/ini.v1 v1.67.0 // indirect
modernc.org/libc v1.66.10 // indirect

4
go.sum
View File

@@ -311,8 +311,8 @@ cloud.google.com/go/gaming v1.6.0/go.mod h1:YMU1GEvA39Qt3zWGyAVA9bpYz/yAhTvaQ1t2
cloud.google.com/go/gaming v1.7.0/go.mod h1:LrB8U7MHdGgFG851iHAfqUdLcKBdQ55hzXy9xBJz0+w=
cloud.google.com/go/gaming v1.8.0/go.mod h1:xAqjS8b7jAVW0KFYeRUxngo9My3f33kFmua++Pi+ggM=
cloud.google.com/go/gaming v1.9.0/go.mod h1:Fc7kEmCObylSWLO334NcO+O9QMDyz+TKC4v1D7X+Bc0=
cloud.google.com/go/geminidataanalytics v0.3.0 h1:2Wi/kqFb5OLuEGH7q+/miE19VTqK1MYHjBEHENap9HI=
cloud.google.com/go/geminidataanalytics v0.3.0/go.mod h1:QRc0b6ywyc3Z7S3etFgslz7hippkW/jRvtops5rKqIg=
cloud.google.com/go/geminidataanalytics v0.5.0 h1:+1usY81Cb+hE8BokpqCM7EgJtRCKzUKx7FvrHbT5hCA=
cloud.google.com/go/geminidataanalytics v0.5.0/go.mod h1:QRc0b6ywyc3Z7S3etFgslz7hippkW/jRvtops5rKqIg=
cloud.google.com/go/gkebackup v0.2.0/go.mod h1:XKvv/4LfG829/B8B7xRkk8zRrOEbKtEam6yNfuQNH60=
cloud.google.com/go/gkebackup v0.3.0/go.mod h1:n/E671i1aOQvUxT541aTkCwExO/bTer2HDlj4TsBRAo=
cloud.google.com/go/gkebackup v0.4.0/go.mod h1:byAyBGUwYGEEww7xsbnUTBHIYcOPy/PgUWUtOeRm9Vg=

View File

@@ -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

View File

@@ -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

View File

@@ -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

View File

@@ -64,14 +64,12 @@ 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
// UserAgentMetadata specifies additional metadata to append to the User-Agent string.
UserAgentMetadata []string
}
type logFormat string
@@ -296,43 +294,6 @@ func (c *ToolConfigs) UnmarshalYAML(ctx context.Context, unmarshal func(interfac
return fmt.Errorf("invalid 'kind' field for tool %q (must be a string)", name)
}
// validify parameter references
if rawParams, ok := v["parameters"]; ok {
if paramsList, ok := rawParams.([]any); ok {
// Turn params into a map
validParamNames := make(map[string]bool)
for _, rawP := range paramsList {
if pMap, ok := rawP.(map[string]any); ok {
if pName, ok := pMap["name"].(string); ok && pName != "" {
validParamNames[pName] = true
}
}
}
// Validate references
for i, rawP := range paramsList {
pMap, ok := rawP.(map[string]any)
if !ok {
continue
}
pName, _ := pMap["name"].(string)
refName, _ := pMap["valueFromParam"].(string)
if refName != "" {
// Check if the referenced parameter exists
if !validParamNames[refName] {
return fmt.Errorf("tool %q config error: parameter %q (index %d) references '%q' in the 'valueFromParam' field, which is not a defined parameter.", name, pName, i, refName)
}
// Check for self-reference
if refName == pName {
return fmt.Errorf("tool %q config error: parameter %q cannot copy value from itself", name, pName)
}
}
}
}
}
yamlDecoder, err := util.NewStrictDecoder(v)
if err != nil {
return fmt.Errorf("error creating YAML decoder for tool %q: %w", name, err)

View File

@@ -27,21 +27,19 @@ import (
v20241105 "github.com/googleapis/genai-toolbox/internal/server/mcp/v20241105"
v20250326 "github.com/googleapis/genai-toolbox/internal/server/mcp/v20250326"
v20250618 "github.com/googleapis/genai-toolbox/internal/server/mcp/v20250618"
v20251125 "github.com/googleapis/genai-toolbox/internal/server/mcp/v20251125"
"github.com/googleapis/genai-toolbox/internal/server/resources"
"github.com/googleapis/genai-toolbox/internal/tools"
)
// LATEST_PROTOCOL_VERSION is the latest version of the MCP protocol supported.
// Update the version used in InitializeResponse when this value is updated.
const LATEST_PROTOCOL_VERSION = v20251125.PROTOCOL_VERSION
const LATEST_PROTOCOL_VERSION = v20250618.PROTOCOL_VERSION
// SUPPORTED_PROTOCOL_VERSIONS is the MCP protocol versions that are supported.
var SUPPORTED_PROTOCOL_VERSIONS = []string{
v20241105.PROTOCOL_VERSION,
v20250326.PROTOCOL_VERSION,
v20250618.PROTOCOL_VERSION,
v20251125.PROTOCOL_VERSION,
}
// InitializeResponse runs capability negotiation and protocol version agreement.
@@ -104,8 +102,6 @@ func NotificationHandler(ctx context.Context, body []byte) error {
// This is the Operation phase of the lifecycle for MCP client-server connections.
func ProcessMethod(ctx context.Context, mcpVersion string, id jsonrpc.RequestId, method string, toolset tools.Toolset, promptset prompts.Promptset, resourceMgr *resources.ResourceManager, body []byte, header http.Header) (any, error) {
switch mcpVersion {
case v20251125.PROTOCOL_VERSION:
return v20251125.ProcessMethod(ctx, id, method, toolset, promptset, resourceMgr, body, header)
case v20250618.PROTOCOL_VERSION:
return v20250618.ProcessMethod(ctx, id, method, toolset, promptset, resourceMgr, body, header)
case v20250326.PROTOCOL_VERSION:

View File

@@ -183,13 +183,6 @@ 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 {

View File

@@ -183,13 +183,6 @@ 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 {

View File

@@ -176,13 +176,6 @@ 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 {

View File

@@ -1,333 +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 v20251125
import (
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"net/http"
"strings"
"github.com/googleapis/genai-toolbox/internal/prompts"
"github.com/googleapis/genai-toolbox/internal/server/mcp/jsonrpc"
"github.com/googleapis/genai-toolbox/internal/server/resources"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/util"
)
// ProcessMethod returns a response for the request.
func ProcessMethod(ctx context.Context, id jsonrpc.RequestId, method string, toolset tools.Toolset, promptset prompts.Promptset, resourceMgr *resources.ResourceManager, body []byte, header http.Header) (any, error) {
switch method {
case PING:
return pingHandler(id)
case TOOLS_LIST:
return toolsListHandler(id, toolset, body)
case TOOLS_CALL:
return toolsCallHandler(ctx, id, resourceMgr, body, header)
case PROMPTS_LIST:
return promptsListHandler(ctx, id, promptset, body)
case PROMPTS_GET:
return promptsGetHandler(ctx, id, resourceMgr, body)
default:
err := fmt.Errorf("invalid method %s", method)
return jsonrpc.NewError(id, jsonrpc.METHOD_NOT_FOUND, err.Error(), nil), err
}
}
// pingHandler handles the "ping" method by returning an empty response.
func pingHandler(id jsonrpc.RequestId) (any, error) {
return jsonrpc.JSONRPCResponse{
Jsonrpc: jsonrpc.JSONRPC_VERSION,
Id: id,
Result: struct{}{},
}, nil
}
func toolsListHandler(id jsonrpc.RequestId, toolset tools.Toolset, body []byte) (any, error) {
var req ListToolsRequest
if err := json.Unmarshal(body, &req); err != nil {
err = fmt.Errorf("invalid mcp tools list request: %w", err)
return jsonrpc.NewError(id, jsonrpc.INVALID_REQUEST, err.Error(), nil), err
}
result := ListToolsResult{
Tools: toolset.McpManifest,
}
return jsonrpc.JSONRPCResponse{
Jsonrpc: jsonrpc.JSONRPC_VERSION,
Id: id,
Result: result,
}, nil
}
// toolsCallHandler generate a response for tools call.
func toolsCallHandler(ctx context.Context, id jsonrpc.RequestId, resourceMgr *resources.ResourceManager, body []byte, header http.Header) (any, error) {
authServices := resourceMgr.GetAuthServiceMap()
// retrieve logger from context
logger, err := util.LoggerFromContext(ctx)
if err != nil {
return jsonrpc.NewError(id, jsonrpc.INTERNAL_ERROR, err.Error(), nil), err
}
var req CallToolRequest
if err = json.Unmarshal(body, &req); err != nil {
err = fmt.Errorf("invalid mcp tools call request: %w", err)
return jsonrpc.NewError(id, jsonrpc.INVALID_REQUEST, err.Error(), nil), err
}
toolName := req.Params.Name
toolArgument := req.Params.Arguments
logger.DebugContext(ctx, fmt.Sprintf("tool name: %s", toolName))
tool, ok := resourceMgr.GetTool(toolName)
if !ok {
err = fmt.Errorf("invalid tool name: tool with name %q does not exist", toolName)
return jsonrpc.NewError(id, jsonrpc.INVALID_PARAMS, err.Error(), nil), err
}
// Get access token
authTokenHeadername, err := tool.GetAuthTokenHeaderName(resourceMgr)
if err != nil {
errMsg := fmt.Errorf("error during invocation: %w", err)
return jsonrpc.NewError(id, jsonrpc.INTERNAL_ERROR, errMsg.Error(), nil), errMsg
}
accessToken := tools.AccessToken(header.Get(authTokenHeadername))
// Check if this specific tool requires the standard authorization header
clientAuth, err := tool.RequiresClientAuthorization(resourceMgr)
if err != nil {
errMsg := fmt.Errorf("error during invocation: %w", err)
return jsonrpc.NewError(id, jsonrpc.INTERNAL_ERROR, errMsg.Error(), nil), errMsg
}
if clientAuth {
if accessToken == "" {
return jsonrpc.NewError(id, jsonrpc.INVALID_REQUEST, "missing access token in the 'Authorization' header", nil), util.ErrUnauthorized
}
}
// marshal arguments and decode it using decodeJSON instead to prevent loss between floats/int.
aMarshal, err := json.Marshal(toolArgument)
if err != nil {
err = fmt.Errorf("unable to marshal tools argument: %w", err)
return jsonrpc.NewError(id, jsonrpc.INTERNAL_ERROR, err.Error(), nil), err
}
var data map[string]any
if err = util.DecodeJSON(bytes.NewBuffer(aMarshal), &data); err != nil {
err = fmt.Errorf("unable to decode tools argument: %w", err)
return jsonrpc.NewError(id, jsonrpc.INTERNAL_ERROR, err.Error(), nil), err
}
// Tool authentication
// claimsFromAuth maps the name of the authservice to the claims retrieved from it.
claimsFromAuth := make(map[string]map[string]any)
// if using stdio, header will be nil and auth will not be supported
if header != nil {
for _, aS := range authServices {
claims, err := aS.GetClaimsFromHeader(ctx, header)
if err != nil {
logger.DebugContext(ctx, err.Error())
continue
}
if claims == nil {
// authService not present in header
continue
}
claimsFromAuth[aS.GetName()] = claims
}
}
// Tool authorization check
verifiedAuthServices := make([]string, len(claimsFromAuth))
i := 0
for k := range claimsFromAuth {
verifiedAuthServices[i] = k
i++
}
// Check if any of the specified auth services is verified
isAuthorized := tool.Authorized(verifiedAuthServices)
if !isAuthorized {
err = fmt.Errorf("unauthorized Tool call: Please make sure your specify correct auth headers: %w", util.ErrUnauthorized)
return jsonrpc.NewError(id, jsonrpc.INVALID_REQUEST, err.Error(), nil), err
}
logger.DebugContext(ctx, "tool invocation authorized")
params, err := tool.ParseParams(data, claimsFromAuth)
if err != nil {
err = fmt.Errorf("provided parameters were invalid: %w", err)
return jsonrpc.NewError(id, jsonrpc.INVALID_PARAMS, err.Error(), nil), err
}
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 {
errStr := err.Error()
// Missing authService tokens.
if errors.Is(err, util.ErrUnauthorized) {
return jsonrpc.NewError(id, jsonrpc.INVALID_REQUEST, err.Error(), nil), err
}
// Upstream auth error
if strings.Contains(errStr, "Error 401") || strings.Contains(errStr, "Error 403") {
if clientAuth {
// Error with client credentials should pass down to the client
return jsonrpc.NewError(id, jsonrpc.INVALID_REQUEST, err.Error(), nil), err
}
// Auth error with ADC should raise internal 500 error
return jsonrpc.NewError(id, jsonrpc.INTERNAL_ERROR, err.Error(), nil), err
}
text := TextContent{
Type: "text",
Text: err.Error(),
}
return jsonrpc.JSONRPCResponse{
Jsonrpc: jsonrpc.JSONRPC_VERSION,
Id: id,
Result: CallToolResult{Content: []TextContent{text}, IsError: true},
}, nil
}
content := make([]TextContent, 0)
sliceRes, ok := results.([]any)
if !ok {
sliceRes = []any{results}
}
for _, d := range sliceRes {
text := TextContent{Type: "text"}
dM, err := json.Marshal(d)
if err != nil {
text.Text = fmt.Sprintf("fail to marshal: %s, result: %s", err, d)
} else {
text.Text = string(dM)
}
content = append(content, text)
}
return jsonrpc.JSONRPCResponse{
Jsonrpc: jsonrpc.JSONRPC_VERSION,
Id: id,
Result: CallToolResult{Content: content},
}, nil
}
// promptsListHandler handles the "prompts/list" method.
func promptsListHandler(ctx context.Context, id jsonrpc.RequestId, promptset prompts.Promptset, body []byte) (any, error) {
// retrieve logger from context
logger, err := util.LoggerFromContext(ctx)
if err != nil {
return jsonrpc.NewError(id, jsonrpc.INTERNAL_ERROR, err.Error(), nil), err
}
logger.DebugContext(ctx, "handling prompts/list request")
var req ListPromptsRequest
if err := json.Unmarshal(body, &req); err != nil {
err = fmt.Errorf("invalid mcp prompts list request: %w", err)
return jsonrpc.NewError(id, jsonrpc.INVALID_REQUEST, err.Error(), nil), err
}
result := ListPromptsResult{
Prompts: promptset.McpManifest,
}
logger.DebugContext(ctx, fmt.Sprintf("returning %d prompts", len(promptset.McpManifest)))
return jsonrpc.JSONRPCResponse{
Jsonrpc: jsonrpc.JSONRPC_VERSION,
Id: id,
Result: result,
}, nil
}
// promptsGetHandler handles the "prompts/get" method.
func promptsGetHandler(ctx context.Context, id jsonrpc.RequestId, resourceMgr *resources.ResourceManager, body []byte) (any, error) {
// retrieve logger from context
logger, err := util.LoggerFromContext(ctx)
if err != nil {
return jsonrpc.NewError(id, jsonrpc.INTERNAL_ERROR, err.Error(), nil), err
}
logger.DebugContext(ctx, "handling prompts/get request")
var req GetPromptRequest
if err := json.Unmarshal(body, &req); err != nil {
err = fmt.Errorf("invalid mcp prompts/get request: %w", err)
return jsonrpc.NewError(id, jsonrpc.INVALID_REQUEST, err.Error(), nil), err
}
promptName := req.Params.Name
logger.DebugContext(ctx, fmt.Sprintf("prompt name: %s", promptName))
prompt, ok := resourceMgr.GetPrompt(promptName)
if !ok {
err := fmt.Errorf("prompt with name %q does not exist", promptName)
return jsonrpc.NewError(id, jsonrpc.INVALID_PARAMS, err.Error(), nil), err
}
// Parse the arguments provided in the request.
argValues, err := prompt.ParseArgs(req.Params.Arguments, nil)
if err != nil {
err = fmt.Errorf("invalid arguments for prompt %q: %w", promptName, err)
return jsonrpc.NewError(id, jsonrpc.INVALID_PARAMS, err.Error(), nil), err
}
logger.DebugContext(ctx, fmt.Sprintf("parsed args: %v", argValues))
// Substitute the argument values into the prompt's messages.
substituted, err := prompt.SubstituteParams(argValues)
if err != nil {
err = fmt.Errorf("error substituting params for prompt %q: %w", promptName, err)
return jsonrpc.NewError(id, jsonrpc.INTERNAL_ERROR, err.Error(), nil), err
}
// Cast the result to the expected []prompts.Message type.
substitutedMessages, ok := substituted.([]prompts.Message)
if !ok {
err = fmt.Errorf("internal error: SubstituteParams returned unexpected type")
return jsonrpc.NewError(id, jsonrpc.INTERNAL_ERROR, err.Error(), nil), err
}
logger.DebugContext(ctx, "substituted params successfully")
// Format the response messages into the required structure.
promptMessages := make([]PromptMessage, len(substitutedMessages))
for i, msg := range substitutedMessages {
promptMessages[i] = PromptMessage{
Role: msg.Role,
Content: TextContent{
Type: "text",
Text: msg.Content,
},
}
}
result := GetPromptResult{
Description: prompt.Manifest().Description,
Messages: promptMessages,
}
return jsonrpc.JSONRPCResponse{
Jsonrpc: jsonrpc.JSONRPC_VERSION,
Id: id,
Result: result,
}, nil
}

View File

@@ -1,219 +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 v20251125
import (
"github.com/googleapis/genai-toolbox/internal/prompts"
"github.com/googleapis/genai-toolbox/internal/server/mcp/jsonrpc"
"github.com/googleapis/genai-toolbox/internal/tools"
)
// SERVER_NAME is the server name used in Implementation.
const SERVER_NAME = "Toolbox"
// PROTOCOL_VERSION is the version of the MCP protocol in this package.
const PROTOCOL_VERSION = "2025-11-25"
// methods that are supported.
const (
PING = "ping"
TOOLS_LIST = "tools/list"
TOOLS_CALL = "tools/call"
PROMPTS_LIST = "prompts/list"
PROMPTS_GET = "prompts/get"
)
/* Empty result */
// EmptyResult represents a response that indicates success but carries no data.
type EmptyResult jsonrpc.Result
/* Pagination */
// Cursor is an opaque token used to represent a cursor for pagination.
type Cursor string
type PaginatedRequest struct {
jsonrpc.Request
Params struct {
// An opaque token representing the current pagination position.
// If provided, the server should return results starting after this cursor.
Cursor Cursor `json:"cursor,omitempty"`
} `json:"params,omitempty"`
}
type PaginatedResult struct {
jsonrpc.Result
// An opaque token representing the pagination position after the last returned result.
// If present, there may be more results available.
NextCursor Cursor `json:"nextCursor,omitempty"`
}
/* Tools */
// Sent from the client to request a list of tools the server has.
type ListToolsRequest struct {
PaginatedRequest
}
// The server's response to a tools/list request from the client.
type ListToolsResult struct {
PaginatedResult
Tools []tools.McpManifest `json:"tools"`
}
// Used by the client to invoke a tool provided by the server.
type CallToolRequest struct {
jsonrpc.Request
Params struct {
Name string `json:"name"`
Arguments map[string]any `json:"arguments,omitempty"`
} `json:"params,omitempty"`
}
// The sender or recipient of messages and data in a conversation.
type Role string
const (
RoleUser Role = "user"
RoleAssistant Role = "assistant"
)
// Base for objects that include optional annotations for the client.
// The client can use annotations to inform how objects are used or displayed
type Annotated struct {
Annotations *struct {
// Describes who the intended customer of this object or data is.
// It can include multiple entries to indicate content useful for multiple
// audiences (e.g., `["user", "assistant"]`).
Audience []Role `json:"audience,omitempty"`
// Describes how important this data is for operating the server.
//
// A value of 1 means "most important," and indicates that the data is
// effectively required, while 0 means "least important," and indicates that
// the data is entirely optional.
//
// @TJS-type number
// @minimum 0
// @maximum 1
Priority float64 `json:"priority,omitempty"`
} `json:"annotations,omitempty"`
}
// TextContent represents text provided to or from an LLM.
type TextContent struct {
Annotated
Type string `json:"type"`
// The text content of the message.
Text string `json:"text"`
}
// The server's response to a tool call.
//
// Any errors that originate from the tool SHOULD be reported inside the result
// object, with `isError` set to true, _not_ as an MCP protocol-level error
// response. Otherwise, the LLM would not be able to see that an error occurred
// and self-correct.
//
// However, any errors in _finding_ the tool, an error indicating that the
// server does not support tool calls, or any other exceptional conditions,
// should be reported as an MCP error response.
type CallToolResult struct {
jsonrpc.Result
// Could be either a TextContent, ImageContent, or EmbeddedResources
// For Toolbox, we will only be sending TextContent
Content []TextContent `json:"content"`
// Whether the tool call ended in an error.
// If not set, this is assumed to be false (the call was successful).
//
// Any errors that originate from the tool SHOULD be reported inside the result
// object, with `isError` set to true, _not_ as an MCP protocol-level error
// response. Otherwise, the LLM would not be able to see that an error occurred
// and self-correct.
//
// However, any errors in _finding_ the tool, an error indicating that the
// server does not support tool calls, or any other exceptional conditions,
// should be reported as an MCP error response.
IsError bool `json:"isError,omitempty"`
// An optional JSON object that represents the structured result of the tool call.
StructuredContent map[string]any `json:"structuredContent,omitempty"`
}
// Additional properties describing a Tool to clients.
//
// NOTE: all properties in ToolAnnotations are **hints**.
// They are not guaranteed to provide a faithful description of
// tool behavior (including descriptive properties like `title`).
//
// Clients should never make tool use decisions based on ToolAnnotations
// received from untrusted servers.
type ToolAnnotations struct {
// A human-readable title for the tool.
Title string `json:"title,omitempty"`
// If true, the tool does not modify its environment.
// Default: false
ReadOnlyHint bool `json:"readOnlyHint,omitempty"`
// If true, the tool may perform destructive updates to its environment.
// If false, the tool performs only additive updates.
// (This property is meaningful only when `readOnlyHint == false`)
// Default: true
DestructiveHint bool `json:"destructiveHint,omitempty"`
// If true, calling the tool repeatedly with the same arguments
// will have no additional effect on the its environment.
// (This property is meaningful only when `readOnlyHint == false`)
// Default: false
IdempotentHint bool `json:"idempotentHint,omitempty"`
// If true, this tool may interact with an "open world" of external
// entities. If false, the tool's domain of interaction is closed.
// For example, the world of a web search tool is open, whereas that
// of a memory tool is not.
// Default: true
OpenWorldHint bool `json:"openWorldHint,omitempty"`
}
/* Prompts */
// Sent from the client to request a list of prompts the server has.
type ListPromptsRequest struct {
PaginatedRequest
}
// The server's response to a prompts/list request from the client.
type ListPromptsResult struct {
PaginatedResult
Prompts []prompts.McpManifest `json:"prompts"`
}
// Used by the client to get a prompt provided by the server.
type GetPromptRequest struct {
jsonrpc.Request
Params struct {
Name string `json:"name"`
Arguments map[string]any `json:"arguments,omitempty"`
} `json:"params"`
}
// The server's response to a prompts/get request from the client.
type GetPromptResult struct {
jsonrpc.Result
Description string `json:"description,omitempty"`
Messages []PromptMessage `json:"messages"`
}
// Describes a message returned as part of a prompt.
type PromptMessage struct {
Role string `json:"role"`
Content TextContent `json:"content"`
}

View File

@@ -37,7 +37,6 @@ const jsonrpcVersion = "2.0"
const protocolVersion20241105 = "2024-11-05"
const protocolVersion20250326 = "2025-03-26"
const protocolVersion20250618 = "2025-06-18"
const protocolVersion20251125 = "2025-11-25"
const serverName = "Toolbox"
var basicInputSchema = map[string]any{
@@ -486,23 +485,6 @@ func TestMcpEndpoint(t *testing.T) {
},
},
},
{
name: "version 2025-11-25",
protocol: protocolVersion20251125,
idHeader: false,
initWant: map[string]any{
"jsonrpc": "2.0",
"id": "mcp-initialize",
"result": map[string]any{
"protocolVersion": "2025-11-25",
"capabilities": map[string]any{
"tools": map[string]any{"listChanged": false},
"prompts": map[string]any{"listChanged": false},
},
"serverInfo": map[string]any{"name": serverName, "version": fakeVersionString},
},
},
},
}
for _, vtc := range versTestCases {
t.Run(vtc.name, func(t *testing.T) {
@@ -512,7 +494,8 @@ func TestMcpEndpoint(t *testing.T) {
if sessionId != "" {
header["Mcp-Session-Id"] = sessionId
}
if vtc.protocol != protocolVersion20241105 && vtc.protocol != protocolVersion20250326 {
if vtc.protocol == protocolVersion20250618 {
header["MCP-Protocol-Version"] = vtc.protocol
}

View File

@@ -158,4 +158,3 @@ func (r *ResourceManager) GetPromptsMap() map[string]prompts.Prompt {
}
return copiedMap
}

View File

@@ -64,11 +64,7 @@ func InitializeConfigs(ctx context.Context, cfg ServerConfig) (
map[string]prompts.Promptset,
error,
) {
metadataStr := cfg.Version
if len(cfg.UserAgentMetadata) > 0 {
metadataStr += "+" + strings.Join(cfg.UserAgentMetadata, "+")
}
ctx = util.WithUserAgent(ctx, metadataStr)
ctx = util.WithUserAgent(ctx, cfg.Version)
instrumentation, err := util.InstrumentationFromContext(ctx)
if err != nil {
panic(err)
@@ -308,14 +304,10 @@ func hostCheck(allowedHosts map[string]struct{}) func(http.Handler) http.Handler
return func(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
_, hasWildcard := allowedHosts["*"]
hostname := r.Host
if host, _, err := net.SplitHostPort(r.Host); err == nil {
hostname = host
}
_, hostIsAllowed := allowedHosts[hostname]
_, hostIsAllowed := allowedHosts[r.Host]
if !hasWildcard && !hostIsAllowed {
// Return 403 Forbidden to block the attack
http.Error(w, "Invalid Host header", http.StatusForbidden)
// Return 400 Bad Request or 403 Forbidden to block the attack
http.Error(w, "Invalid Host header", http.StatusBadRequest)
return
}
next.ServeHTTP(w, r)
@@ -414,11 +406,7 @@ func NewServer(ctx context.Context, cfg ServerConfig) (*Server, error) {
}
allowedHostsMap := make(map[string]struct{}, len(cfg.AllowedHosts))
for _, h := range cfg.AllowedHosts {
hostname := h
if host, _, err := net.SplitHostPort(h); err == nil {
hostname = host
}
allowedHostsMap[hostname] = struct{}{}
allowedHostsMap[h] = struct{}{}
}
r.Use(hostCheck(allowedHostsMap))

View File

@@ -14,23 +14,23 @@
package cloudgda
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net/http"
geminidataanalytics "cloud.google.com/go/geminidataanalytics/apiv1beta"
"cloud.google.com/go/geminidataanalytics/apiv1beta/geminidataanalyticspb"
"github.com/goccy/go-yaml"
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/util"
"go.opentelemetry.io/otel/trace"
"golang.org/x/oauth2"
"golang.org/x/oauth2/google"
"google.golang.org/api/option"
)
const SourceKind string = "cloud-gemini-data-analytics"
const Endpoint string = "https://geminidataanalytics.googleapis.com"
// NewDataChatClient can be overridden for testing.
var NewDataChatClient = geminidataanalytics.NewDataChatClient
// validate interface
var _ sources.SourceConfig = Config{}
@@ -67,29 +67,19 @@ func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.So
return nil, fmt.Errorf("error in User Agent retrieval: %s", err)
}
var client *http.Client
if r.UseClientOAuth {
client = &http.Client{
Transport: util.NewUserAgentRoundTripper(ua, http.DefaultTransport),
}
} else {
// Use Application Default Credentials
// Scope: "https://www.googleapis.com/auth/cloud-platform" is generally sufficient for GDA
creds, err := google.FindDefaultCredentials(ctx, "https://www.googleapis.com/auth/cloud-platform")
if err != nil {
return nil, fmt.Errorf("failed to find default credentials: %w", err)
}
baseClient := oauth2.NewClient(ctx, creds.TokenSource)
baseClient.Transport = util.NewUserAgentRoundTripper(ua, baseClient.Transport)
client = baseClient
}
s := &Source{
Config: r,
Client: client,
BaseURL: Endpoint,
userAgent: ua,
}
if !r.UseClientOAuth {
client, err := NewDataChatClient(ctx, option.WithUserAgent(ua))
if err != nil {
return nil, fmt.Errorf("failed to create DataChatClient: %w", err)
}
s.Client = client
}
return s, nil
}
@@ -97,8 +87,7 @@ var _ sources.Source = &Source{}
type Source struct {
Config
Client *http.Client
BaseURL string
Client *geminidataanalytics.DataChatClient
userAgent string
}
@@ -114,63 +103,36 @@ func (s *Source) GetProjectID() string {
return s.ProjectID
}
func (s *Source) GetBaseURL() string {
return s.BaseURL
}
func (s *Source) GetClient(ctx context.Context, accessToken string) (*http.Client, error) {
if s.UseClientOAuth {
if accessToken == "" {
return nil, fmt.Errorf("client-side OAuth is enabled but no access token was provided")
}
token := &oauth2.Token{AccessToken: accessToken}
baseClient := oauth2.NewClient(ctx, oauth2.StaticTokenSource(token))
baseClient.Transport = util.NewUserAgentRoundTripper(s.userAgent, baseClient.Transport)
return baseClient, nil
}
return s.Client, nil
}
func (s *Source) UseClientAuthorization() bool {
return s.UseClientOAuth
}
func (s *Source) RunQuery(ctx context.Context, tokenStr string, bodyBytes []byte) (any, error) {
// The API endpoint itself always uses the "global" location.
apiLocation := "global"
apiParent := fmt.Sprintf("projects/%s/locations/%s", s.GetProjectID(), apiLocation)
apiURL := fmt.Sprintf("%s/v1beta/%s:queryData", s.GetBaseURL(), apiParent)
func (s *Source) GetClient(ctx context.Context, tokenStr string) (*geminidataanalytics.DataChatClient, func(), error) {
if s.UseClientOAuth {
if tokenStr == "" {
return nil, nil, fmt.Errorf("client-side OAuth is enabled but no access token was provided")
}
token := &oauth2.Token{AccessToken: tokenStr}
opts := []option.ClientOption{
option.WithUserAgent(s.userAgent),
option.WithTokenSource(oauth2.StaticTokenSource(token)),
}
client, err := s.GetClient(ctx, tokenStr)
if err != nil {
return nil, fmt.Errorf("failed to get HTTP client: %w", err)
client, err := NewDataChatClient(ctx, opts...)
if err != nil {
return nil, nil, fmt.Errorf("failed to create per-request DataChatClient: %w", err)
}
return client, func() { client.Close() }, nil
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, apiURL, bytes.NewBuffer(bodyBytes))
if err != nil {
return nil, fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
resp, err := client.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to execute request: %w", err)
}
defer resp.Body.Close()
respBody, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("failed to read response body: %w", err)
}
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(respBody))
}
var result map[string]any
if err := json.Unmarshal(respBody, &result); err != nil {
return nil, fmt.Errorf("failed to unmarshal response: %w", err)
}
return result, nil
return s.Client, func() {}, nil
}
func (s *Source) RunQuery(ctx context.Context, tokenStr string, req *geminidataanalyticspb.QueryDataRequest) (*geminidataanalyticspb.QueryDataResponse, error) {
client, cleanup, err := s.GetClient(ctx, tokenStr)
if err != nil {
return nil, err
}
defer cleanup()
return client.QueryData(ctx, req)
}

View File

@@ -181,11 +181,9 @@ func TestInitialize(t *testing.T) {
if gdaSrc.Client == nil && !tc.wantClientOAuth {
t.Fatal("expected non-nil HTTP client for ADC, got nil")
}
// When client OAuth is true, the source's client should be initialized with a base HTTP client
// that includes the user agent round tripper, but not the OAuth token. The token-aware
// client is created by GetClient.
if gdaSrc.Client == nil && tc.wantClientOAuth {
t.Fatal("expected non-nil HTTP client for client OAuth config, got nil")
// When client OAuth is true, the source's client should be nil.
if gdaSrc.Client != nil && tc.wantClientOAuth {
t.Fatal("expected nil HTTP client for client OAuth config, got non-nil")
}
// Test UseClientAuthorization method
@@ -195,15 +193,16 @@ func TestInitialize(t *testing.T) {
// Test GetClient with accessToken for client OAuth scenarios
if tc.wantClientOAuth {
client, err := gdaSrc.GetClient(ctx, "dummy-token")
client, cleanup, err := gdaSrc.GetClient(ctx, "dummy-token")
if err != nil {
t.Fatalf("GetClient with token failed: %v", err)
}
defer cleanup()
if client == nil {
t.Fatal("expected non-nil HTTP client from GetClient with token, got nil")
}
// Ensure passing empty token with UseClientOAuth enabled returns error
_, err = gdaSrc.GetClient(ctx, "")
_, _, err = gdaSrc.GetClient(ctx, "")
if err == nil || err.Error() != "client-side OAuth is enabled but no access token was provided" {
t.Errorf("expected 'client-side OAuth is enabled but no access token was provided' error, got: %v", err)
}

View File

@@ -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" {

View File

@@ -19,11 +19,13 @@ import (
"encoding/json"
"fmt"
"cloud.google.com/go/geminidataanalytics/apiv1beta/geminidataanalyticspb"
"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/protobuf/encoding/protojson"
)
const kind string = "cloud-gemini-data-analytics-query"
@@ -60,7 +62,49 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.T
type compatibleSource interface {
GetProjectID() string
UseClientAuthorization() bool
RunQuery(context.Context, string, []byte) (any, error)
RunQuery(context.Context, string, *geminidataanalyticspb.QueryDataRequest) (*geminidataanalyticspb.QueryDataResponse, error)
}
// QueryDataContext wraps geminidataanalyticspb.QueryDataContext to support YAML decoding via protojson.
type QueryDataContext struct {
*geminidataanalyticspb.QueryDataContext
}
func (q *QueryDataContext) UnmarshalYAML(b []byte) error {
var raw map[string]any
if err := yaml.Unmarshal(b, &raw); err != nil {
return fmt.Errorf("failed to unmarshal context from yaml: %w", err)
}
jsonBytes, err := json.Marshal(raw)
if err != nil {
return fmt.Errorf("failed to marshal context map: %w", err)
}
q.QueryDataContext = &geminidataanalyticspb.QueryDataContext{}
if err := protojson.Unmarshal(jsonBytes, q.QueryDataContext); err != nil {
return fmt.Errorf("failed to unmarshal context to proto: %w", err)
}
return nil
}
// GenerationOptions wraps geminidataanalyticspb.GenerationOptions to support YAML decoding via protojson.
type GenerationOptions struct {
*geminidataanalyticspb.GenerationOptions
}
func (g *GenerationOptions) UnmarshalYAML(b []byte) error {
var raw map[string]any
if err := yaml.Unmarshal(b, &raw); err != nil {
return fmt.Errorf("failed to unmarshal generation options from yaml: %w", err)
}
jsonBytes, err := json.Marshal(raw)
if err != nil {
return fmt.Errorf("failed to marshal generation options map: %w", err)
}
g.GenerationOptions = &geminidataanalyticspb.GenerationOptions{}
if err := protojson.Unmarshal(jsonBytes, g.GenerationOptions); err != nil {
return fmt.Errorf("failed to unmarshal generation options to proto: %w", err)
}
return nil
}
type Config struct {
@@ -97,12 +141,14 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error)
}
mcpManifest := tools.GetMcpManifest(cfg.Name, cfg.Description, cfg.AuthRequired, allParameters, nil)
return Tool{
t := Tool{
Config: cfg,
AllParams: allParameters,
manifest: tools.Manifest{Description: cfg.Description, Parameters: allParameters.Manifest(), AuthRequired: cfg.AuthRequired},
mcpManifest: mcpManifest,
}, nil
}
return t, nil
}
// validate interface
@@ -145,18 +191,20 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
// The parent in the request payload uses the tool's configured location.
payloadParent := fmt.Sprintf("projects/%s/locations/%s", source.GetProjectID(), t.Location)
payload := &QueryDataRequest{
Parent: payloadParent,
Prompt: query,
Context: t.Context,
GenerationOptions: t.GenerationOptions,
req := &geminidataanalyticspb.QueryDataRequest{
Parent: payloadParent,
Prompt: query,
}
bodyBytes, err := json.Marshal(payload)
if err != nil {
return nil, fmt.Errorf("failed to marshal request payload: %w", err)
if t.Context != nil {
req.Context = t.Context.QueryDataContext
}
return source.RunQuery(ctx, tokenStr, bodyBytes)
if t.GenerationOptions != nil {
req.GenerationOptions = t.GenerationOptions.GenerationOptions
}
return source.RunQuery(ctx, tokenStr, req)
}
func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (parameters.ParamValues, error) {

View File

@@ -16,19 +16,16 @@ package cloudgda_test
import (
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"net/http/httptest"
"testing"
"cloud.google.com/go/geminidataanalytics/apiv1beta/geminidataanalyticspb"
yaml "github.com/goccy/go-yaml"
"github.com/google/go-cmp/cmp"
"github.com/google/go-cmp/cmp/cmpopts"
"github.com/googleapis/genai-toolbox/internal/server"
"github.com/googleapis/genai-toolbox/internal/server/resources"
"github.com/googleapis/genai-toolbox/internal/sources"
cloudgdasrc "github.com/googleapis/genai-toolbox/internal/sources/cloudgda"
"github.com/googleapis/genai-toolbox/internal/testutils"
"github.com/googleapis/genai-toolbox/internal/tools"
cloudgdatool "github.com/googleapis/genai-toolbox/internal/tools/cloudgda"
@@ -74,23 +71,29 @@ func TestParseFromYaml(t *testing.T) {
Location: "us-central1",
AuthRequired: []string{},
Context: &cloudgdatool.QueryDataContext{
DatasourceReferences: &cloudgdatool.DatasourceReferences{
SpannerReference: &cloudgdatool.SpannerReference{
DatabaseReference: &cloudgdatool.SpannerDatabaseReference{
ProjectID: "cloud-db-nl2sql",
Region: "us-central1",
InstanceID: "evalbench",
DatabaseID: "financial",
Engine: cloudgdatool.SpannerEngineGoogleSQL,
},
AgentContextReference: &cloudgdatool.AgentContextReference{
ContextSetID: "projects/cloud-db-nl2sql/locations/us-east1/contextSets/bdf_gsql_gemini_all_templates",
QueryDataContext: &geminidataanalyticspb.QueryDataContext{
DatasourceReferences: &geminidataanalyticspb.DatasourceReferences{
References: &geminidataanalyticspb.DatasourceReferences_SpannerReference{
SpannerReference: &geminidataanalyticspb.SpannerReference{
DatabaseReference: &geminidataanalyticspb.SpannerDatabaseReference{
ProjectId: "cloud-db-nl2sql",
Region: "us-central1",
InstanceId: "evalbench",
DatabaseId: "financial",
Engine: geminidataanalyticspb.SpannerDatabaseReference_GOOGLE_SQL,
},
AgentContextReference: &geminidataanalyticspb.AgentContextReference{
ContextSetId: "projects/cloud-db-nl2sql/locations/us-east1/contextSets/bdf_gsql_gemini_all_templates",
},
},
},
},
},
},
GenerationOptions: &cloudgdatool.GenerationOptions{
GenerateQueryResult: true,
GenerationOptions: &geminidataanalyticspb.GenerationOptions{
GenerateQueryResult: true,
},
},
},
},
@@ -108,68 +111,63 @@ func TestParseFromYaml(t *testing.T) {
if err != nil {
t.Fatalf("unable to unmarshal: %s", err)
}
if !cmp.Equal(tc.want, got.Tools) {
t.Fatalf("incorrect parse: want %v, got %v", tc.want, got.Tools)
if !cmp.Equal(tc.want, got.Tools, cmpopts.IgnoreUnexported(geminidataanalyticspb.QueryDataContext{}, geminidataanalyticspb.DatasourceReferences{}, geminidataanalyticspb.SpannerReference{}, geminidataanalyticspb.SpannerDatabaseReference{}, geminidataanalyticspb.AgentContextReference{}, geminidataanalyticspb.GenerationOptions{}, geminidataanalyticspb.DatasourceReferences_SpannerReference{})) {
t.Errorf("incorrect parse: want %v, got %v", tc.want, got.Tools)
}
})
}
}
// authRoundTripper is a mock http.RoundTripper that adds a dummy Authorization header.
type authRoundTripper struct {
Token string
Next http.RoundTripper
// fakeSource implements the compatibleSource interface for testing.
type fakeSource struct {
projectID string
useClientOAuth bool
expectedQuery string
expectedParent string
response *geminidataanalyticspb.QueryDataResponse
}
func (rt *authRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
newReq := *req
newReq.Header = make(http.Header)
for k, v := range req.Header {
newReq.Header[k] = v
}
newReq.Header.Set("Authorization", rt.Token)
if rt.Next == nil {
return http.DefaultTransport.RoundTrip(&newReq)
}
return rt.Next.RoundTrip(&newReq)
func (f *fakeSource) GetProjectID() string {
return f.projectID
}
type mockSource struct {
kind string
client *http.Client // Can be used to inject a specific client
baseURL string // BaseURL is needed to implement sources.Source.BaseURL
config cloudgdasrc.Config // to return from ToConfig
func (f *fakeSource) UseClientAuthorization() bool {
return f.useClientOAuth
}
func (m *mockSource) SourceKind() string { return m.kind }
func (m *mockSource) ToConfig() sources.SourceConfig { return m.config }
func (m *mockSource) GetClient(ctx context.Context, token string) (*http.Client, error) {
if m.client != nil {
return m.client, nil
}
// Default client for testing if not explicitly set
transport := &http.Transport{}
authTransport := &authRoundTripper{
Token: "Bearer test-access-token", // Dummy token
Next: transport,
}
return &http.Client{Transport: authTransport}, nil
func (f *fakeSource) SourceKind() string {
return "fake-gda-source"
}
func (m *mockSource) UseClientAuthorization() bool { return false }
func (m *mockSource) Initialize(ctx context.Context, tracer interface{}) (sources.Source, error) {
return m, nil
func (f *fakeSource) ToConfig() sources.SourceConfig {
return nil
}
func (f *fakeSource) Initialize(ctx context.Context, tracer interface{}) (sources.Source, error) {
return f, nil
}
func (f *fakeSource) RunQuery(ctx context.Context, token string, req *geminidataanalyticspb.QueryDataRequest) (*geminidataanalyticspb.QueryDataResponse, error) {
if req.Prompt != f.expectedQuery {
return nil, fmt.Errorf("unexpected query: got %q, want %q", req.Prompt, f.expectedQuery)
}
if req.Parent != f.expectedParent {
return nil, fmt.Errorf("unexpected parent: got %q, want %q", req.Parent, f.expectedParent)
}
// Basic validation of context/options could be added here if needed,
// but the test case mainly checks if they are passed correctly via successful invocation.
return f.response, nil
}
func (m *mockSource) BaseURL() string { return m.baseURL }
func TestInitialize(t *testing.T) {
t.Parallel()
// Minimal fake source
fake := &fakeSource{projectID: "test-project"}
srcs := map[string]sources.Source{
"gda-api-source": &cloudgdasrc.Source{
Config: cloudgdasrc.Config{Name: "gda-api-source", Kind: cloudgdasrc.SourceKind, ProjectID: "test-project"},
Client: &http.Client{},
BaseURL: cloudgdasrc.Endpoint,
},
"gda-api-source": fake,
}
tcs := []struct {
@@ -188,9 +186,6 @@ func TestInitialize(t *testing.T) {
},
}
// Add an incompatible source for testing
srcs["incompatible-source"] = &mockSource{kind: "another-kind"}
for _, tc := range tcs {
tc := tc
t.Run(tc.desc, func(t *testing.T) {
@@ -207,92 +202,27 @@ func TestInitialize(t *testing.T) {
func TestInvoke(t *testing.T) {
t.Parallel()
// Mock the HTTP client and server for Invoke testing
serverMux := http.NewServeMux()
// Update expected URL path to include the location "us-central1"
serverMux.HandleFunc("/v1beta/projects/test-project/locations/global:queryData", func(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodPost {
t.Errorf("expected POST method, got %s", r.Method)
http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
return
}
if r.Header.Get("Content-Type") != "application/json" {
t.Errorf("expected Content-Type application/json, got %s", r.Header.Get("Content-Type"))
http.Error(w, "Bad request", http.StatusBadRequest)
return
}
// Read and unmarshal the request body
bodyBytes, err := io.ReadAll(r.Body)
if err != nil {
t.Errorf("failed to read request body: %v", err)
http.Error(w, "Internal Server Error", http.StatusInternalServerError)
return
}
var reqPayload cloudgdatool.QueryDataRequest
if err := json.Unmarshal(bodyBytes, &reqPayload); err != nil {
t.Errorf("failed to unmarshal request payload: %v", err)
http.Error(w, "Bad request", http.StatusBadRequest)
return
}
projectID := "test-project"
location := "us-central1"
query := "How many accounts who have region in Prague are eligible for loans?"
expectedParent := fmt.Sprintf("projects/%s/locations/%s", projectID, location)
// Verify expected fields
if r.Header.Get("Authorization") == "" {
t.Errorf("expected Authorization header, got empty")
http.Error(w, "Unauthorized", http.StatusUnauthorized)
return
}
if reqPayload.Prompt != "How many accounts who have region in Prague are eligible for loans?" {
t.Errorf("unexpected prompt: %s", reqPayload.Prompt)
}
// Verify payload's parent uses the tool's configured location
if reqPayload.Parent != fmt.Sprintf("projects/%s/locations/%s", "test-project", "us-central1") {
t.Errorf("unexpected payload parent: got %q, want %q", reqPayload.Parent, fmt.Sprintf("projects/%s/locations/%s", "test-project", "us-central1"))
}
// Verify context from config
if reqPayload.Context == nil ||
reqPayload.Context.DatasourceReferences == nil ||
reqPayload.Context.DatasourceReferences.SpannerReference == nil ||
reqPayload.Context.DatasourceReferences.SpannerReference.DatabaseReference == nil ||
reqPayload.Context.DatasourceReferences.SpannerReference.DatabaseReference.ProjectID != "cloud-db-nl2sql" {
t.Errorf("unexpected context: %v", reqPayload.Context)
}
// Verify generation options from config
if reqPayload.GenerationOptions == nil || !reqPayload.GenerationOptions.GenerateQueryResult {
t.Errorf("unexpected generation options: %v", reqPayload.GenerationOptions)
}
// Simulate a successful response
resp := map[string]any{
"queryResult": "SELECT count(*) FROM accounts WHERE region = 'Prague' AND eligible_for_loans = true;",
"naturalLanguageAnswer": "There are 5 accounts in Prague eligible for loans.",
}
_ = json.NewEncoder(w).Encode(resp)
})
mockServer := httptest.NewServer(serverMux)
defer mockServer.Close()
ctx := testutils.ContextWithUserAgent(context.Background(), "test-user-agent")
// Create an authenticated client that uses the mock server
authTransport := &authRoundTripper{
Token: "Bearer test-access-token",
Next: mockServer.Client().Transport,
// Prepare expected response
expectedResp := &geminidataanalyticspb.QueryDataResponse{
GeneratedQuery: "SELECT count(*) FROM accounts WHERE region = 'Prague' AND eligible_for_loans = true;",
NaturalLanguageAnswer: "There are 5 accounts in Prague eligible for loans.",
}
authClient := &http.Client{Transport: authTransport}
// Create a real cloudgdasrc.Source but inject the authenticated client
mockGdaSource := &cloudgdasrc.Source{
Config: cloudgdasrc.Config{Name: "mock-gda-source", Kind: cloudgdasrc.SourceKind, ProjectID: "test-project"},
Client: authClient,
BaseURL: mockServer.URL,
fake := &fakeSource{
projectID: projectID,
expectedQuery: query,
expectedParent: expectedParent,
response: expectedResp,
}
srcs := map[string]sources.Source{
"mock-gda-source": mockGdaSource,
"mock-gda-source": fake,
}
// Initialize the tool config with context
@@ -301,25 +231,31 @@ func TestInvoke(t *testing.T) {
Kind: "cloud-gemini-data-analytics-query",
Source: "mock-gda-source",
Description: "Query Gemini Data Analytics",
Location: "us-central1", // Set location for the test
Location: location,
Context: &cloudgdatool.QueryDataContext{
DatasourceReferences: &cloudgdatool.DatasourceReferences{
SpannerReference: &cloudgdatool.SpannerReference{
DatabaseReference: &cloudgdatool.SpannerDatabaseReference{
ProjectID: "cloud-db-nl2sql",
Region: "us-central1",
InstanceID: "evalbench",
DatabaseID: "financial",
Engine: cloudgdatool.SpannerEngineGoogleSQL,
},
AgentContextReference: &cloudgdatool.AgentContextReference{
ContextSetID: "projects/cloud-db-nl2sql/locations/us-east1/contextSets/bdf_gsql_gemini_all_templates",
QueryDataContext: &geminidataanalyticspb.QueryDataContext{
DatasourceReferences: &geminidataanalyticspb.DatasourceReferences{
References: &geminidataanalyticspb.DatasourceReferences_SpannerReference{
SpannerReference: &geminidataanalyticspb.SpannerReference{
DatabaseReference: &geminidataanalyticspb.SpannerDatabaseReference{
ProjectId: "cloud-db-nl2sql",
Region: "us-central1",
InstanceId: "evalbench",
DatabaseId: "financial",
Engine: geminidataanalyticspb.SpannerDatabaseReference_GOOGLE_SQL,
},
AgentContextReference: &geminidataanalyticspb.AgentContextReference{
ContextSetId: "projects/cloud-db-nl2sql/locations/us-east1/contextSets/bdf_gsql_gemini_all_templates",
},
},
},
},
},
},
GenerationOptions: &cloudgdatool.GenerationOptions{
GenerateQueryResult: true,
GenerationOptions: &geminidataanalyticspb.GenerationOptions{
GenerateQueryResult: true,
},
},
}
@@ -330,24 +266,25 @@ func TestInvoke(t *testing.T) {
// Prepare parameters for invocation - ONLY query
params := parameters.ParamValues{
{Name: "query", Value: "How many accounts who have region in Prague are eligible for loans?"},
{Name: "query", Value: query},
}
resourceMgr := resources.NewResourceManager(srcs, nil, nil, nil, nil, nil, nil)
ctx := testutils.ContextWithUserAgent(context.Background(), "test-user-agent")
// Invoke the tool
result, err := tool.Invoke(ctx, resourceMgr, params, "") // No accessToken needed for ADC client
result, err := tool.Invoke(ctx, resourceMgr, params, "")
if err != nil {
t.Fatalf("tool invocation failed: %v", err)
}
// Validate the result
expectedResult := map[string]any{
"queryResult": "SELECT count(*) FROM accounts WHERE region = 'Prague' AND eligible_for_loans = true;",
"naturalLanguageAnswer": "There are 5 accounts in Prague eligible for loans.",
gotResp, ok := result.(*geminidataanalyticspb.QueryDataResponse)
if !ok {
t.Fatalf("expected result type *geminidataanalyticspb.QueryDataResponse, got %T", result)
}
if !cmp.Equal(expectedResult, result) {
t.Errorf("unexpected result: got %v, want %v", result, expectedResult)
if diff := cmp.Diff(expectedResp, gotResp, cmpopts.IgnoreUnexported(geminidataanalyticspb.QueryDataResponse{})); diff != "" {
t.Errorf("unexpected result mismatch (-want +got):\n%s", diff)
}
}

View File

@@ -1,116 +0,0 @@
// Copyright 2025 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 cloudgda
// See full service definition at: https://github.com/googleapis/googleapis/blob/master/google/cloud/geminidataanalytics/v1beta/data_chat_service.proto
// QueryDataRequest represents the JSON body for the queryData API
type QueryDataRequest struct {
Parent string `json:"parent"`
Prompt string `json:"prompt"`
Context *QueryDataContext `json:"context,omitempty"`
GenerationOptions *GenerationOptions `json:"generationOptions,omitempty"`
}
// QueryDataContext reflects the proto definition for the query context.
type QueryDataContext struct {
DatasourceReferences *DatasourceReferences `json:"datasourceReferences,omitempty" yaml:"datasourceReferences,omitempty"`
}
// DatasourceReferences reflects the proto definition for datasource references, using a oneof.
type DatasourceReferences struct {
SpannerReference *SpannerReference `json:"spannerReference,omitempty" yaml:"spannerReference,omitempty"`
AlloyDBReference *AlloyDBReference `json:"alloydb,omitempty" yaml:"alloydb,omitempty"`
CloudSQLReference *CloudSQLReference `json:"cloudSqlReference,omitempty" yaml:"cloudSqlReference,omitempty"`
}
// SpannerReference reflects the proto definition for Spanner database reference.
type SpannerReference struct {
DatabaseReference *SpannerDatabaseReference `json:"databaseReference,omitempty" yaml:"databaseReference,omitempty"`
AgentContextReference *AgentContextReference `json:"agentContextReference,omitempty" yaml:"agentContextReference,omitempty"`
}
// SpannerDatabaseReference reflects the proto definition for a Spanner database reference.
type SpannerDatabaseReference struct {
Engine SpannerEngine `json:"engine,omitempty" yaml:"engine,omitempty"`
ProjectID string `json:"projectId,omitempty" yaml:"projectId,omitempty"`
Region string `json:"region,omitempty" yaml:"region,omitempty"`
InstanceID string `json:"instanceId,omitempty" yaml:"instanceId,omitempty"`
DatabaseID string `json:"databaseId,omitempty" yaml:"databaseId,omitempty"`
TableIDs []string `json:"tableIds,omitempty" yaml:"tableIds,omitempty"`
}
// SpannerEngine represents the engine of the Spanner instance.
type SpannerEngine string
const (
SpannerEngineUnspecified SpannerEngine = "ENGINE_UNSPECIFIED"
SpannerEngineGoogleSQL SpannerEngine = "GOOGLE_SQL"
SpannerEnginePostgreSQL SpannerEngine = "POSTGRESQL"
)
// AlloyDBReference reflects the proto definition for an AlloyDB database reference.
type AlloyDBReference struct {
DatabaseReference *AlloyDBDatabaseReference `json:"databaseReference,omitempty" yaml:"databaseReference,omitempty"`
AgentContextReference *AgentContextReference `json:"agentContextReference,omitempty" yaml:"agentContextReference,omitempty"`
}
// AlloyDBDatabaseReference reflects the proto definition for an AlloyDB database reference.
type AlloyDBDatabaseReference struct {
ProjectID string `json:"projectId,omitempty" yaml:"projectId,omitempty"`
Region string `json:"region,omitempty" yaml:"region,omitempty"`
ClusterID string `json:"clusterId,omitempty" yaml:"clusterId,omitempty"`
InstanceID string `json:"instanceId,omitempty" yaml:"instanceId,omitempty"`
DatabaseID string `json:"databaseId,omitempty" yaml:"databaseId,omitempty"`
TableIDs []string `json:"tableIds,omitempty" yaml:"tableIds,omitempty"`
}
// CloudSQLReference reflects the proto definition for a Cloud SQL database reference.
type CloudSQLReference struct {
DatabaseReference *CloudSQLDatabaseReference `json:"databaseReference,omitempty" yaml:"databaseReference,omitempty"`
AgentContextReference *AgentContextReference `json:"agentContextReference,omitempty" yaml:"agentContextReference,omitempty"`
}
// CloudSQLDatabaseReference reflects the proto definition for a Cloud SQL database reference.
type CloudSQLDatabaseReference struct {
Engine CloudSQLEngine `json:"engine,omitempty" yaml:"engine,omitempty"`
ProjectID string `json:"projectId,omitempty" yaml:"projectId,omitempty"`
Region string `json:"region,omitempty" yaml:"region,omitempty"`
InstanceID string `json:"instanceId,omitempty" yaml:"instanceId,omitempty"`
DatabaseID string `json:"databaseId,omitempty" yaml:"databaseId,omitempty"`
TableIDs []string `json:"tableIds,omitempty" yaml:"tableIds,omitempty"`
}
// CloudSQLEngine represents the engine of the Cloud SQL instance.
type CloudSQLEngine string
const (
CloudSQLEngineUnspecified CloudSQLEngine = "ENGINE_UNSPECIFIED"
CloudSQLEnginePostgreSQL CloudSQLEngine = "POSTGRESQL"
CloudSQLEngineMySQL CloudSQLEngine = "MYSQL"
)
// AgentContextReference reflects the proto definition for agent context.
type AgentContextReference struct {
ContextSetID string `json:"contextSetId,omitempty" yaml:"contextSetId,omitempty"`
}
// GenerationOptions reflects the proto definition for generation options.
type GenerationOptions struct {
GenerateQueryResult bool `json:"generateQueryResult" yaml:"generateQueryResult"`
GenerateNaturalLanguageAnswer bool `json:"generateNaturalLanguageAnswer" yaml:"generateNaturalLanguageAnswer"`
GenerateExplanation bool `json:"generateExplanation" yaml:"generateExplanation"`
GenerateDisambiguationQuestion bool `json:"generateDisambiguationQuestion" yaml:"generateDisambiguationQuestion"`
}

View File

@@ -103,12 +103,10 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
if !ok {
return nil, fmt.Errorf("invalid or missing '%s' parameter; expected a string", pageURLKey)
}
var tokenStr string
if source.UseClientAuthorization() {
tokenStr, err = accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
tokenStr, err := accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
return source.FHIRFetchPage(ctx, url, tokenStr)
}

View File

@@ -131,12 +131,9 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
return nil, fmt.Errorf("invalid or missing '%s' parameter; expected a string", patientIDKey)
}
var tokenStr string
if source.UseClientAuthorization() {
tokenStr, err = accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
tokenStr, err := accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
var opts []googleapi.CallOption

View File

@@ -161,12 +161,9 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
return nil, err
}
var tokenStr string
if source.UseClientAuthorization() {
tokenStr, err = accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
tokenStr, err := accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
var summary bool

View File

@@ -95,12 +95,9 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
if err != nil {
return nil, err
}
var tokenStr string
if source.UseClientAuthorization() {
tokenStr, err = accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
tokenStr, err := accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
return source.GetDataset(tokenStr)
}

View File

@@ -116,12 +116,9 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
if err != nil {
return nil, err
}
var tokenStr string
if source.UseClientAuthorization() {
tokenStr, err = accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
tokenStr, err := accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
return source.GetDICOMStore(storeID, tokenStr)
}

View File

@@ -116,12 +116,9 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
if err != nil {
return nil, err
}
var tokenStr string
if source.UseClientAuthorization() {
tokenStr, err = accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
tokenStr, err := accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
return source.GetDICOMStoreMetrics(storeID, tokenStr)
}

View File

@@ -130,12 +130,9 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
if !ok {
return nil, fmt.Errorf("invalid or missing '%s' parameter; expected a string", idKey)
}
var tokenStr string
if source.UseClientAuthorization() {
tokenStr, err = accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
tokenStr, err := accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
return source.GetFHIRResource(storeID, resType, resID, tokenStr)
}

View File

@@ -116,12 +116,9 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
if err != nil {
return nil, err
}
var tokenStr string
if source.UseClientAuthorization() {
tokenStr, err = accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
tokenStr, err := accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
return source.GetFHIRStore(storeID, tokenStr)
}

View File

@@ -116,12 +116,9 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
if err != nil {
return nil, err
}
var tokenStr string
if source.UseClientAuthorization() {
tokenStr, err = accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
tokenStr, err := accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
return source.GetFHIRStoreMetrics(storeID, tokenStr)
}

View File

@@ -95,12 +95,9 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
if err != nil {
return nil, err
}
var tokenStr string
if source.UseClientAuthorization() {
tokenStr, err = accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
tokenStr, err := accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
return source.ListDICOMStores(tokenStr)
}

View File

@@ -95,12 +95,9 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
if err != nil {
return nil, err
}
var tokenStr string
if source.UseClientAuthorization() {
tokenStr, err = accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
tokenStr, err := accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
return source.ListFHIRStores(tokenStr)
}

View File

@@ -127,12 +127,9 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
if err != nil {
return nil, err
}
var tokenStr string
if source.UseClientAuthorization() {
tokenStr, err = accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
tokenStr, err := accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
study, ok := params.AsMap()[studyInstanceUIDKey].(string)
if !ok {

View File

@@ -140,12 +140,9 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
if err != nil {
return nil, err
}
var tokenStr string
if source.UseClientAuthorization() {
tokenStr, err = accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
tokenStr, err := accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
opts, err := common.ParseDICOMSearchParameters(params, []string{sopInstanceUIDKey, patientNameKey, patientIDKey, accessionNumberKey, referringPhysicianNameKey, studyDateKey, modalityKey})

View File

@@ -138,12 +138,9 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
if err != nil {
return nil, err
}
var tokenStr string
if source.UseClientAuthorization() {
tokenStr, err = accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
tokenStr, err := accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
opts, err := common.ParseDICOMSearchParameters(params, []string{seriesInstanceUIDKey, patientNameKey, patientIDKey, accessionNumberKey, referringPhysicianNameKey, studyDateKey, modalityKey})

View File

@@ -133,12 +133,9 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
if err != nil {
return nil, err
}
var tokenStr string
if source.UseClientAuthorization() {
tokenStr, err = accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
tokenStr, err := accessToken.ParseBearerToken()
if err != nil {
return nil, fmt.Errorf("error parsing access token: %w", err)
}
opts, err := common.ParseDICOMSearchParameters(params, []string{studyInstanceUIDKey, patientNameKey, patientIDKey, accessionNumberKey, referringPhysicianNameKey, studyDateKey})
if err != nil {

View File

@@ -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
}

View File

@@ -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)
}
})
}
}

View File

@@ -134,12 +134,7 @@ func ParseParams(ps Parameters, data map[string]any, claimsMap map[string]map[st
var err error
paramAuthServices := p.GetAuthServices()
name := p.GetName()
sourceParamName := p.GetValueFromParam()
if sourceParamName != "" {
v, _ = data[sourceParamName]
} else if len(paramAuthServices) == 0 {
if len(paramAuthServices) == 0 {
// parse non auth-required parameter
var ok bool
v, ok = data[name]
@@ -323,7 +318,6 @@ type Parameter interface {
GetRequired() bool
GetAuthServices() []ParamAuthService
GetEmbeddedBy() string
GetValueFromParam() string
Parse(any) (any, error)
Manifest() ParameterManifest
McpManifest() (ParameterMcpManifest, []string)
@@ -471,9 +465,6 @@ func ParseParameter(ctx context.Context, p map[string]any, paramType string) (Pa
func (ps Parameters) Manifest() []ParameterManifest {
rtn := make([]ParameterManifest, 0, len(ps))
for _, p := range ps {
if p.GetValueFromParam() != "" {
continue
}
rtn = append(rtn, p.Manifest())
}
return rtn
@@ -485,11 +476,6 @@ func (ps Parameters) McpManifest() (McpToolsSchema, map[string][]string) {
authParam := make(map[string][]string)
for _, p := range ps {
// If the parameter is sourced from another param, skip it in the MCP manifest
if p.GetValueFromParam() != "" {
continue
}
name := p.GetName()
paramManifest, authParamList := p.McpManifest()
defaultV := p.GetDefault()
@@ -523,7 +509,6 @@ type ParameterManifest struct {
Default any `json:"default,omitempty"`
AdditionalProperties any `json:"additionalProperties,omitempty"`
EmbeddedBy string `json:"embeddedBy,omitempty"`
ValueFromParam string `json:"valueFromParam,omitempty"`
}
// ParameterMcpManifest represents properties when served as part of a ToolMcpManifest.
@@ -546,7 +531,6 @@ type CommonParameter struct {
AuthServices []ParamAuthService `yaml:"authServices"`
AuthSources []ParamAuthService `yaml:"authSources"` // Deprecated: Kept for compatibility.
EmbeddedBy string `yaml:"embeddedBy"`
ValueFromParam string `yaml:"valueFromParam"`
}
// GetName returns the name specified for the Parameter.
@@ -604,16 +588,10 @@ func (p *CommonParameter) IsExcludedValues(v any) bool {
return false
}
// GetEmbeddedBy returns the embedding model name for the Parameter.
func (p *CommonParameter) GetEmbeddedBy() string {
return p.EmbeddedBy
}
// GetValueFromParam returns the param value to copy from.
func (p *CommonParameter) GetValueFromParam() string {
return p.ValueFromParam
}
// MatchStringOrRegex checks if the input matches the target
func MatchStringOrRegex(input, target any) bool {
targetS, ok := target.(string)

View File

@@ -31,18 +31,6 @@
"default": "tools.yaml",
"isRequired": false
},
{
"type": "named",
"name": "--tools-files",
"description": "Multiple file paths specifying tool configurations. Files will be merged. Cannot be used with -tools-file or -tools-folder.",
"isRequired": false
},
{
"type": "named",
"name": "--tools-folder",
"description": "Directory path containing YAML tool configuration files. All .yaml and .yml files in the directory will be loaded and merged. Cannot be used with -tools-file or -tools-files.",
"isRequired": false
},
{
"type": "named",
"name": "--address",
@@ -82,82 +70,6 @@
"warn",
"error"
]
},
{
"type": "named",
"name": "--logging-format",
"description": "Specify logging format to use.",
"default": "standard",
"choices": [
"standard",
"json"
]
},
{
"type": "named",
"name": "--disable-reload",
"description": "Disables dynamic reloading of tools file.",
"format": "boolean",
"isRequired": false
},
{
"type": "named",
"name": "--prebuilt",
"description": "Use a prebuilt tool configuration by source type.",
"isRequired": false
},
{
"type": "named",
"name": "--stdio",
"description": "Listens via MCP STDIO instead of acting as a remote HTTP server.",
"format": "boolean",
"isRequired": false
},
{
"type": "named",
"name": "--telemetry-gcp",
"description": "Enable exporting directly to Google Cloud Monitoring.",
"format": "boolean",
"isRequired": false
},
{
"type": "named",
"name": "--telemetry-otlp",
"description": "Enable exporting using OpenTelemetry Protocol (OTLP) to the specified endpoint (e.g. 'http://127.0.0.1:4318').",
"isRequired": false
},
{
"type": "named",
"name": "--telemetry-service-name",
"description": "Sets the value of the service.name resource attribute for telemetry data.",
"default": "toolbox",
"isRequired": false
},
{
"type": "named",
"name": "--ui",
"description": "Launches the Toolbox UI web server.",
"format": "boolean",
"isRequired": false
},
{
"type": "named",
"name": "--allowed-origins",
"description": "Specifies a list of origins permitted to access this server.",
"default": "*",
"isRequired": false
},
{
"type": "named",
"name": "--help",
"description": "Show help for toolbox",
"isRequired": false
},
{
"type": "named",
"name": "--version",
"description": "Show version for toolbox",
"isRequired": false
}
]
}

View File

@@ -147,20 +147,12 @@ 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...)

View File

@@ -18,78 +18,75 @@ import (
"bytes"
"context"
"encoding/json"
"fmt"
"net"
"net/http"
"net/http/httptest"
"net/url"
"regexp"
"strings"
"testing"
"time"
geminidataanalytics "cloud.google.com/go/geminidataanalytics/apiv1beta"
"cloud.google.com/go/geminidataanalytics/apiv1beta/geminidataanalyticspb"
"github.com/googleapis/genai-toolbox/internal/server/mcp/jsonrpc"
source "github.com/googleapis/genai-toolbox/internal/sources/cloudgda"
"github.com/googleapis/genai-toolbox/internal/testutils"
"github.com/googleapis/genai-toolbox/internal/tools/cloudgda"
"github.com/googleapis/genai-toolbox/tests"
"google.golang.org/api/option"
"google.golang.org/grpc"
"google.golang.org/grpc/credentials/insecure"
)
var (
cloudGdaToolKind = "cloud-gemini-data-analytics-query"
)
type cloudGdaTransport struct {
transport http.RoundTripper
url *url.URL
}
func (t *cloudGdaTransport) RoundTrip(req *http.Request) (*http.Response, error) {
if strings.HasPrefix(req.URL.String(), "https://geminidataanalytics.googleapis.com") {
req.URL.Scheme = t.url.Scheme
req.URL.Host = t.url.Host
}
return t.transport.RoundTrip(req)
}
type masterHandler struct {
type mockDataChatServer struct {
geminidataanalyticspb.UnimplementedDataChatServiceServer
t *testing.T
}
func (h *masterHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
if !strings.Contains(r.UserAgent(), "genai-toolbox/") {
h.t.Errorf("User-Agent header not found")
func (s *mockDataChatServer) QueryData(ctx context.Context, req *geminidataanalyticspb.QueryDataRequest) (*geminidataanalyticspb.QueryDataResponse, error) {
if req.Prompt == "" {
s.t.Errorf("missing prompt")
return nil, fmt.Errorf("missing prompt")
}
if r.Method != http.MethodPost {
http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
return
}
return &geminidataanalyticspb.QueryDataResponse{
GeneratedQuery: "SELECT * FROM table;",
NaturalLanguageAnswer: "Here is the answer.",
}, nil
}
// Verify URL structure
// Expected: /v1beta/projects/{project}/locations/global:queryData
if !strings.Contains(r.URL.Path, ":queryData") || !strings.Contains(r.URL.Path, "locations/global") {
h.t.Errorf("unexpected URL path: %s", r.URL.Path)
http.Error(w, "Not found", http.StatusNotFound)
return
}
var reqBody cloudgda.QueryDataRequest
if err := json.NewDecoder(r.Body).Decode(&reqBody); err != nil {
h.t.Fatalf("failed to decode request body: %v", err)
}
if reqBody.Prompt == "" {
http.Error(w, "missing prompt", http.StatusBadRequest)
return
}
response := map[string]any{
"queryResult": "SELECT * FROM table;",
"naturalLanguageAnswer": "Here is the answer.",
}
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusOK)
if err := json.NewEncoder(w).Encode(response); err != nil {
http.Error(w, err.Error(), http.StatusInternalServerError)
func getCloudGdaToolsConfig() map[string]any {
return map[string]any{
"sources": map[string]any{
"my-gda-source": map[string]any{
"kind": "cloud-gemini-data-analytics",
"projectId": "test-project",
},
},
"tools": map[string]any{
"cloud-gda-query": map[string]any{
"kind": cloudGdaToolKind,
"source": "my-gda-source",
"description": "Test GDA Tool",
"location": "us-central1",
"context": map[string]any{
"datasourceReferences": map[string]any{
"spannerReference": map[string]any{
"databaseReference": map[string]any{
"projectId": "test-project",
"instanceId": "test-instance",
"databaseId": "test-db",
"engine": "GOOGLE_SQL",
},
},
},
},
},
},
}
}
@@ -97,26 +94,37 @@ func TestCloudGdaToolEndpoints(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
defer cancel()
handler := &masterHandler{t: t}
server := httptest.NewServer(handler)
defer server.Close()
serverURL, err := url.Parse(server.URL)
// Start a gRPC server
lis, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatalf("failed to parse server URL: %v", err)
t.Fatalf("failed to listen: %v", err)
}
s := grpc.NewServer()
geminidataanalyticspb.RegisterDataChatServiceServer(s, &mockDataChatServer{t: t})
go func() {
if err := s.Serve(lis); err != nil {
// This might happen on strict shutdown, log if unexpected
t.Logf("server executed: %v", err)
}
}()
defer s.Stop()
originalTransport := http.DefaultClient.Transport
if originalTransport == nil {
originalTransport = http.DefaultTransport
// Configure toolbox to use the gRPC server
endpoint := lis.Addr().String()
// Override client creation
origFunc := source.NewDataChatClient
defer func() {
source.NewDataChatClient = origFunc
}()
source.NewDataChatClient = func(ctx context.Context, opts ...option.ClientOption) (*geminidataanalytics.DataChatClient, error) {
opts = append(opts,
option.WithEndpoint(endpoint),
option.WithoutAuthentication(),
option.WithGRPCDialOption(grpc.WithTransportCredentials(insecure.NewCredentials())))
return origFunc(ctx, opts...)
}
http.DefaultClient.Transport = &cloudGdaTransport{
transport: originalTransport,
url: serverURL,
}
t.Cleanup(func() {
http.DefaultClient.Transport = originalTransport
})
var args []string
toolsFile := getCloudGdaToolsConfig()
@@ -156,7 +164,7 @@ func TestCloudGdaToolEndpoints(t *testing.T) {
// 2. RunToolInvokeParametersTest
params := []byte(`{"query": "test question"}`)
tests.RunToolInvokeParametersTest(t, toolName, params, "\"queryResult\":\"SELECT * FROM table;\"")
tests.RunToolInvokeParametersTest(t, toolName, params, "\"generated_query\":\"SELECT * FROM table;\"")
// 3. Manual MCP Tool Call Test
// Initialize MCP session
@@ -196,38 +204,3 @@ func TestCloudGdaToolEndpoints(t *testing.T) {
t.Errorf("MCP response does not contain expected query result: %s", respStr)
}
}
func getCloudGdaToolsConfig() map[string]any {
// Mocked responses and a dummy `projectId` are used in this integration
// test due to limited project-specific allowlisting. API functionality is
// verified via internal monitoring; this test specifically validates the
// integration flow between the source and the tool.
return map[string]any{
"sources": map[string]any{
"my-gda-source": map[string]any{
"kind": "cloud-gemini-data-analytics",
"projectId": "test-project",
},
},
"tools": map[string]any{
"cloud-gda-query": map[string]any{
"kind": cloudGdaToolKind,
"source": "my-gda-source",
"description": "Test GDA Tool",
"location": "us-central1",
"context": map[string]any{
"datasourceReferences": map[string]any{
"spannerReference": map[string]any{
"databaseReference": map[string]any{
"projectId": "test-project",
"instanceId": "test-instance",
"databaseId": "test-db",
"engine": "GOOGLE_SQL",
},
},
},
},
},
},
}
}

View File

@@ -112,7 +112,8 @@ func TestHealthcareToolEndpoints(t *testing.T) {
fhirStoreID := "fhir-store-" + uuid.New().String()
dicomStoreID := "dicom-store-" + uuid.New().String()
patient1ID, patient2ID := setupHealthcareResources(t, healthcareService, healthcareDataset, fhirStoreID, dicomStoreID)
patient1ID, patient2ID, teardown := setupHealthcareResources(t, healthcareService, healthcareDataset, fhirStoreID, dicomStoreID)
defer teardown(t)
toolsFile := getToolsConfig(sourceConfig)
toolsFile = addClientAuthSourceConfig(t, toolsFile)
@@ -172,8 +173,10 @@ func TestHealthcareToolWithStoreRestriction(t *testing.T) {
disallowedFHIRStoreID := "fhir-store-disallowed-" + uuid.New().String()
disallowedDICOMStoreID := "dicom-store-disallowed-" + uuid.New().String()
setupHealthcareResources(t, healthcareService, healthcareDataset, allowedFHIRStoreID, allowedDICOMStoreID)
setupHealthcareResources(t, healthcareService, healthcareDataset, disallowedFHIRStoreID, disallowedDICOMStoreID)
_, _, teardownAllowedStores := setupHealthcareResources(t, healthcareService, healthcareDataset, allowedFHIRStoreID, allowedDICOMStoreID)
defer teardownAllowedStores(t)
_, _, teardownDisallowedStores := setupHealthcareResources(t, healthcareService, healthcareDataset, disallowedFHIRStoreID, disallowedDICOMStoreID)
defer teardownDisallowedStores(t)
// Configure source with dataset restriction.
sourceConfig["allowedFhirStores"] = []string{allowedFHIRStoreID}
@@ -254,7 +257,7 @@ func newHealthcareService(ctx context.Context) (*healthcare.Service, error) {
return healthcareService, nil
}
func setupHealthcareResources(t *testing.T, service *healthcare.Service, datasetID, fhirStoreID, dicomStoreID string) (string, string) {
func setupHealthcareResources(t *testing.T, service *healthcare.Service, datasetID, fhirStoreID, dicomStoreID string) (string, string, func(*testing.T)) {
datasetName := fmt.Sprintf("projects/%s/locations/%s/datasets/%s", healthcareProject, healthcareRegion, datasetID)
var err error
@@ -263,24 +266,12 @@ func setupHealthcareResources(t *testing.T, service *healthcare.Service, dataset
if fhirStore, err = service.Projects.Locations.Datasets.FhirStores.Create(datasetName, fhirStore).FhirStoreId(fhirStoreID).Do(); err != nil {
t.Fatalf("failed to create fhir store: %v", err)
}
// Register cleanup
t.Cleanup(func() {
if _, err := service.Projects.Locations.Datasets.FhirStores.Delete(fhirStore.Name).Do(); err != nil {
t.Logf("failed to delete fhir store: %v", err)
}
})
// Create DICOM store
dicomStore := &healthcare.DicomStore{}
if dicomStore, err = service.Projects.Locations.Datasets.DicomStores.Create(datasetName, dicomStore).DicomStoreId(dicomStoreID).Do(); err != nil {
t.Fatalf("failed to create dicom store: %v", err)
}
// Register cleanup
t.Cleanup(func() {
if _, err := service.Projects.Locations.Datasets.DicomStores.Delete(dicomStore.Name).Do(); err != nil {
t.Logf("failed to delete dicom store: %v", err)
}
})
// Create Patient 1
patient1Body := bytes.NewBuffer([]byte(`{
@@ -326,7 +317,15 @@ func setupHealthcareResources(t *testing.T, service *healthcare.Service, dataset
createFHIRResource(t, service, fhirStore.Name, "Observation", observation2Body)
}
return patient1ID, patient2ID
teardown := func(t *testing.T) {
if _, err := service.Projects.Locations.Datasets.FhirStores.Delete(fhirStore.Name).Do(); err != nil {
t.Logf("failed to delete fhir store: %v", err)
}
if _, err := service.Projects.Locations.Datasets.DicomStores.Delete(dicomStore.Name).Do(); err != nil {
t.Logf("failed to delete dicom store: %v", err)
}
}
return patient1ID, patient2ID, teardown
}
func getToolsConfig(sourceConfig map[string]any) map[string]any {

View File

@@ -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",
},
},
}
}

View File

@@ -132,20 +132,12 @@ 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 {
@@ -194,7 +186,6 @@ 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

View File

@@ -1,251 +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 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",
"valueFromParam": "content",
},
},
}
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"`,
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",
},
},
},
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
}

View File

@@ -111,10 +111,6 @@ 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")
@@ -122,10 +118,6 @@ 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)
@@ -173,5 +165,4 @@ 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")
}

View File

@@ -1240,10 +1240,7 @@ func RunPostgresListTablesTest(t *testing.T, tableNameParam, tableNameAuth, user
var filteredGot []any
for _, item := range got {
if tableMap, ok := item.(map[string]interface{}); ok {
name, _ := tableMap["object_name"].(string)
// Only keep the table if it matches expected test tables
if name == tableNameParam || name == tableNameAuth {
if schema, ok := tableMap["schema_name"]; ok && schema == "public" {
filteredGot = append(filteredGot, item)
}
}