feat(serverless-spark): add list/get sessions commands (googleapis#2576)

## Description

Adds commands for Serverless Spark sessions analogous to existing
list/get commands for batches.

Unlike batches, we don't plan to add create session commands, as these
are generally only created via the Jupyter protocol. Agents will need to
have other ways to create sessions, for example by interacting directly
with an IDE like Jupyter Lab or VS Code, but those won't be part of MCP
Toolbox.

## PR Checklist

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

- [ ] Make sure you reviewed

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

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

🛠️ Part of googleapis#2405

Co-authored-by: Wenxin Du <117315983+duwenxin99@users.noreply.github.com>

Author: dborowitz

cmd/internal/imports.go

@@ -204,7 +204,9 @@ import (
 	_ "github.com/googleapis/genai-toolbox/internal/tools/serverlessspark/serverlesssparkcreatepysparkbatch"
 	_ "github.com/googleapis/genai-toolbox/internal/tools/serverlessspark/serverlesssparkcreatesparkbatch"
 	_ "github.com/googleapis/genai-toolbox/internal/tools/serverlessspark/serverlesssparkgetbatch"
+	_ "github.com/googleapis/genai-toolbox/internal/tools/serverlessspark/serverlesssparkgetsession"
 	_ "github.com/googleapis/genai-toolbox/internal/tools/serverlessspark/serverlesssparklistbatches"
+	_ "github.com/googleapis/genai-toolbox/internal/tools/serverlessspark/serverlesssparklistsessions"
 	_ "github.com/googleapis/genai-toolbox/internal/tools/singlestore/singlestoreexecutesql"
 	_ "github.com/googleapis/genai-toolbox/internal/tools/singlestore/singlestoresql"
 	_ "github.com/googleapis/genai-toolbox/internal/tools/snowflake/snowflakeexecutesql"

cmd/internal/tools_file_test.go

@@ -1753,7 +1753,7 @@ func TestPrebuiltTools(t *testing.T) {
 			wantToolset: server.ToolsetConfigs{
 				"serverless_spark_tools": tools.ToolsetConfig{
 					Name:      "serverless_spark_tools",
-					ToolNames: []string{"list_batches", "get_batch", "cancel_batch", "create_pyspark_batch", "create_spark_batch"},
+					ToolNames: []string{"list_batches", "get_batch", "cancel_batch", "create_pyspark_batch", "create_spark_batch", "list_sessions", "get_session"},
 				},
 			},
 		},

docs/en/reference/prebuilt-tools.md

@@ -746,6 +746,8 @@ See [Usage Examples](../reference/cli.md#examples).
     *   `cancel_batch`: Cancels a Spark batch.
     *   `create_pyspark_batch`: Creates a PySpark batch.
     *   `create_spark_batch`: Creates a Spark batch.
+    *   `list_sessions`: Lists Spark sessions.
+    *   `get_session`: Gets a Spark session.
 
 ## SingleStore
 

docs/en/resources/sources/serverless-spark.md

@@ -25,6 +25,10 @@ Apache Spark.
   Create a Serverless Spark PySpark batch operation.
 - [`serverless-spark-create-spark-batch`](../tools/serverless-spark/serverless-spark-create-spark-batch.md)
   Create a Serverless Spark Java batch operation.
+- [`serverless-spark-list-sessions`](../tools/serverless-spark/serverless-spark-list-sessions.md)
+  List and filter Serverless Spark sessions.
+- [`serverless-spark-get-session`](../tools/serverless-spark/serverless-spark-get-session.md)
+  Get a Serverless Spark session.
 
 ## Requirements
 

docs/en/resources/tools/serverless-spark/_index.md

@@ -11,3 +11,5 @@ description: >
 - [serverless-spark-cancel-batch](./serverless-spark-cancel-batch.md)
 - [serverless-spark-create-pyspark-batch](./serverless-spark-create-pyspark-batch.md)
 - [serverless-spark-create-spark-batch](./serverless-spark-create-spark-batch.md)
+- [serverless-spark-get-session](./serverless-spark-get-session.md)
+- [serverless-spark-list-sessions](./serverless-spark-list-sessions.md)

docs/en/resources/tools/serverless-spark/serverless-spark-get-session.md

@@ -0,0 +1,58 @@
+---
+title: "serverless-spark-get-session"
+type: docs
+weight: 1
+description: >
+  A "serverless-spark-get-session" tool retrieves a specific Spark session from the source.
+aliases:
+  - /resources/tools/serverless-spark-get-session
+---
+
+## About
+
+A `serverless-spark-get-session` tool retrieves a specific Spark session from a
+Google Cloud Serverless for Apache Spark source. It's compatible with the
+following sources:
+
+- [serverless-spark](../../sources/serverless-spark.md)
+
+`serverless-spark-get-session` accepts the following parameters:
+
+- **`name`** (required): The short name of the session, e.g. for `projects/my-project/locations/us-central1/sessions/my-session`, pass `my-session`.
+
+The tool gets the `project` and `location` from the source configuration.
+
+## Example
+
+```yaml
+kind: tools
+name: get_spark_session
+type: serverless-spark-get-session
+source: my-serverless-spark-source
+description: Use this tool to get details of a serverless spark session.
+```
+
+## Response Format
+
+```json
+{
+  "consoleUrl": "https://console.cloud.google.com/dataproc/interactive/us-central1/my-session/details?project=my-project",
+  "logsUrl": "https://console.cloud.google.com/logs/viewer?...",
+  "session": {  
+    "name": "projects/my-project/locations/us-central1/sessions/my-session",
+    "uuid": "a1b2c3d4-e5f6-7890-1234-567890abcdef",
+    "state": "ACTIVE",
+    // ... complete session resource definition
+  }
+}
+```
+
+## Reference
+
+| **field**    | **type** | **required** | **description**                                    |
+| ------------ | :------: | :----------: | -------------------------------------------------- |
+| type         |  string  |     true     | Must be "serverless-spark-get-session".            |
+| source       |  string  |     true     | Name of the source the tool should use.            |
+| description  |  string  |     true     | Description of the tool that is passed to the LLM. |
+| authRequired | string[] |    false     | List of auth services required to invoke this tool |
+  

docs/en/resources/tools/serverless-spark/serverless-spark-list-sessions.md

@@ -0,0 +1,82 @@
+---
+title: "serverless-spark-list-sessions"
+type: docs
+weight: 1
+description: >
+  A "serverless-spark-list-sessions" tool returns a list of Spark sessions from the source.
+aliases:
+  - /resources/tools/serverless-spark-list-sessions
+---
+
+## About
+
+A `serverless-spark-list-sessions` tool returns a list of Spark sessions from a
+Google Cloud Serverless for Apache Spark source. It's compatible with the
+following sources:
+
+- [serverless-spark](../../sources/serverless-spark.md)
+
+`serverless-spark-list-sessions` accepts the following parameters:
+
+- **`filter`** (optional): Optional. A filter for the sessions to return in the
+  response. A filter is a logical expression constraining the values of various
+  fields in each session resource. Filters are case sensitive, and may contain
+  multiple clauses combined with logical operators (AND, OR). Supported fields
+  are session_id, session_uuid, state, create_time, and labels. Example: `state
+  = ACTIVE and create_time < "2023-01-01T00:00:00Z"` is a filter for sessions in
+  an ACTIVE state that were created before 2023-01-01. `state = ACTIVE and
+  labels.environment=production` is a filter for sessions in an ACTIVE state
+  that have a production environment label.
+- **`pageSize`** (optional): The maximum number of sessions to return in a single
+  page. Defaults to `20`.
+- **`pageToken`** (optional): A page token, received from a previous call, to
+  retrieve the next page of results.
+
+The tool gets the `project` and `location` from the source configuration.
+
+## Example
+
+```yaml
+kind: tools
+name: list_spark_sessions
+type: serverless-spark-list-sessions
+source: my-serverless-spark-source
+description: Use this tool to list and filter serverless spark sessions.
+```
+
+## Response Format
+
+```json
+{
+  "sessions": [
+    {
+      "name": "projects/my-project/locations/us-central1/sessions/session-abc-123",
+      "uuid": "a1b2c3d4-e5f6-7890-1234-567890abcdef",
+      "state": "ACTIVE",
+      "creator": "alice@example.com",
+      "createTime": "2023-10-27T10:00:00Z",
+      "consoleUrl": "https://console.cloud.google.com/dataproc/interactive/us-central1/session-abc-123/details?project=my-project",
+      "logsUrl": "https://console.cloud.google.com/logs/viewer?..."
+    },
+    {
+      "name": "projects/my-project/locations/us-central1/sessions/session-def-456",
+      "uuid": "b2c3d4e5-f6a7-8901-2345-678901bcdefa",
+      "state": "TERMINATED",
+      "creator": "alice@example.com",
+      "createTime": "2023-10-27T11:30:00Z",
+      "consoleUrl": "https://console.cloud.google.com/dataproc/interactive/us-central1/session-def-456/details?project=my-project",
+      "logsUrl": "https://console.cloud.google.com/logs/viewer?..."
+    }
+  ],
+  "nextPageToken": "abcd1234"
+}
+```
+
+## Reference
+
+| **field**    | **type** | **required** | **description**                                    |
+| ------------ | :------: | :----------: | -------------------------------------------------- |
+| type         |  string  |     true     | Must be "serverless-spark-list-sessions".          |
+| source       |  string  |     true     | Name of the source the tool should use.            |
+| description  |  string  |     true     | Description of the tool that is passed to the LLM. |
+| authRequired | string[] |    false     | List of auth services required to invoke this tool |

internal/prebuiltconfigs/tools/serverless-spark.yaml

@@ -34,6 +34,12 @@ tools:
   create_spark_batch:
     kind: serverless-spark-create-spark-batch
     source: serverless-spark-source
+  list_sessions:
+    kind: serverless-spark-list-sessions
+    source: serverless-spark-source
+  get_session:
+    kind: serverless-spark-get-session
+    source: serverless-spark-source
 
 toolsets:
   serverless_spark_tools:
@@ -42,3 +48,5 @@ toolsets:
     - cancel_batch
     - create_pyspark_batch
     - create_spark_batch
+    - list_sessions
+    - get_session

internal/sources/serverlessspark/serverlessspark.go

@@ -17,6 +17,7 @@ package serverlessspark
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"time"
 
@@ -77,11 +78,16 @@ func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.So
 	if err != nil {
 		return nil, fmt.Errorf("failed to create longrunning client: %w", err)
 	}
+	sessionClient, err := dataproc.NewSessionControllerClient(ctx, option.WithEndpoint(endpoint), option.WithUserAgent(ua))
+	if err != nil {
+		return nil, fmt.Errorf("failed to create dataproc session client: %w", err)
+	}
 
 	s := &Source{
-		Config:    r,
-		Client:    client,
-		OpsClient: opsClient,
+		Config:        r,
+		Client:        client,
+		OpsClient:     opsClient,
+		SessionClient: sessionClient,
 	}
 	return s, nil
 }
@@ -90,8 +96,9 @@ var _ sources.Source = &Source{}
 
 type Source struct {
 	Config
-	Client    *dataproc.BatchControllerClient
-	OpsClient *longrunning.OperationsClient
+	Client        *dataproc.BatchControllerClient
+	OpsClient     *longrunning.OperationsClient
+	SessionClient *dataproc.SessionControllerClient
 }
 
 func (s *Source) SourceType() string {
@@ -114,18 +121,16 @@ func (s *Source) GetBatchControllerClient() *dataproc.BatchControllerClient {
 	return s.Client
 }
 
+func (s *Source) GetSessionControllerClient() *dataproc.SessionControllerClient {
+	return s.SessionClient
+}
+
 func (s *Source) GetOperationsClient(ctx context.Context) (*longrunning.OperationsClient, error) {
 	return s.OpsClient, nil
 }
 
 func (s *Source) Close() error {
-	if err := s.Client.Close(); err != nil {
-		return err
-	}
-	if err := s.OpsClient.Close(); err != nil {
-		return err
-	}
-	return nil
+	return errors.Join(s.Client.Close(), s.SessionClient.Close(), s.OpsClient.Close())
 }
 
 func (s *Source) CancelOperation(ctx context.Context, operation string) (any, error) {
@@ -292,3 +297,119 @@ func (s *Source) GetBatch(ctx context.Context, name string) (map[string]any, err
 
 	return wrappedResult, nil
 }
+
+// ListSessionsResponse is the response from the list sessions API.
+type ListSessionsResponse struct {
+	Sessions      []Session `json:"sessions"`
+	NextPageToken string    `json:"nextPageToken"`
+}
+
+// Session represents a single session job.
+type Session struct {
+	Name       string `json:"name"`
+	UUID       string `json:"uuid"`
+	State      string `json:"state"`
+	Creator    string `json:"creator"`
+	CreateTime string `json:"createTime"`
+	ConsoleURL string `json:"consoleUrl"`
+	LogsURL    string `json:"logsUrl"`
+}
+
+func (s *Source) ListSessions(ctx context.Context, ps *int, pt, filter string) (any, error) {
+	client := s.GetSessionControllerClient()
+	parent := fmt.Sprintf("projects/%s/locations/%s", s.GetProject(), s.GetLocation())
+	req := &dataprocpb.ListSessionsRequest{
+		Parent: parent,
+	}
+
+	if ps != nil {
+		req.PageSize = int32(*ps)
+	}
+	if pt != "" {
+		req.PageToken = pt
+	}
+	if filter != "" {
+		req.Filter = filter
+	}
+
+	it := client.ListSessions(ctx, req)
+	pager := iterator.NewPager(it, int(req.PageSize), req.PageToken)
+
+	var sessionPbs []*dataprocpb.Session
+	nextPageToken, err := pager.NextPage(&sessionPbs)
+	if err != nil {
+		return nil, fmt.Errorf("failed to list sessions: %w", err)
+	}
+
+	sessions, err := ToSessions(sessionPbs)
+	if err != nil {
+		return nil, err
+	}
+
+	return ListSessionsResponse{Sessions: sessions, NextPageToken: nextPageToken}, nil
+}
+
+func (s *Source) GetSession(ctx context.Context, name string) (map[string]any, error) {
+	client := s.GetSessionControllerClient()
+	req := &dataprocpb.GetSessionRequest{
+		Name: fmt.Sprintf("projects/%s/locations/%s/sessions/%s", s.GetProject(), s.GetLocation(), name),
+	}
+
+	sessionPb, err := client.GetSession(ctx, req)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get session: %w", err)
+	}
+
+	jsonBytes, err := protojson.Marshal(sessionPb)
+	if err != nil {
+		return nil, fmt.Errorf("failed to marshal session to JSON: %w", err)
+	}
+
+	var result map[string]any
+	if err := json.Unmarshal(jsonBytes, &result); err != nil {
+		return nil, fmt.Errorf("failed to unmarshal session JSON: %w", err)
+	}
+
+	consoleUrl, err := SessionConsoleURLFromProto(sessionPb)
+	if err != nil {
+		return nil, fmt.Errorf("error generating console url: %v", err)
+	}
+	logsUrl, err := SessionLogsURLFromProto(sessionPb)
+	if err != nil {
+		return nil, fmt.Errorf("error generating logs url: %v", err)
+	}
+
+	wrappedResult := map[string]any{
+		"consoleUrl": consoleUrl,
+		"logsUrl":    logsUrl,
+		"session":    result,
+	}
+
+	return wrappedResult, nil
+}
+
+// ToSessions converts a slice of protobuf Session messages to a slice of Session structs.
+func ToSessions(sessionPbs []*dataprocpb.Session) ([]Session, error) {
+	sessions := make([]Session, 0, len(sessionPbs))
+	for _, sessionPb := range sessionPbs {
+		consoleUrl, err := SessionConsoleURLFromProto(sessionPb)
+		if err != nil {
+			return nil, fmt.Errorf("error generating console url: %v", err)
+		}
+		logsUrl, err := SessionLogsURLFromProto(sessionPb)
+		if err != nil {
+			return nil, fmt.Errorf("error generating logs url: %v", err)
+		}
+		session := Session{
+			Name:       sessionPb.Name,
+			UUID:       sessionPb.Uuid,
+			State:      sessionPb.State.Enum().String(),
+			Creator:    sessionPb.Creator,
+			CreateTime: sessionPb.CreateTime.AsTime().Format(time.RFC3339),
+			ConsoleURL: consoleUrl,
+			LogsURL:    logsUrl,
+		}
+		sessions = append(sessions, session)
+	}
+	return sessions, nil
+}