enhance connection handling

Author: adranwit

README.md

@@ -65,7 +65,7 @@ The easiest way to try the toolbox is to run the server with both HTTP and
 STDIO transports enabled:
 
 
-```bash
+```bash cd
 go run ./cmd/mcp-sqlkit -a :5000 
 ```
 
@@ -185,16 +185,15 @@ The toolbox registers the following MCP tools (see `mcp/tool.go`).
 | `dbQuery`              | Execute a SQL query and return the result set | `db/query.Input`            |
 | `dbExec`               | Execute DML/DDL and return rows affected      | `db/exec.Input`             |
 | `dbListConnections`    | List connectors visible to the caller         | `db/connector.ListInput`    |
-| `dbAddConnection`      | Register a new connector                      | `db/connector.ConnectionInput` |
-| `dbUpdateConnection`   | Update an existing connector (upsert)         | `db/connector.ConnectionInput` |
+| `dbSetConnection`      | Create or update a connector (upsert)         | `db/connector.ConnectionInput` |
 | `dbListTables`         | List tables for a given catalog/schema        | `db/meta.ListTablesInput`   |
 | `dbListColumns`        | List columns for a given table                | `db/meta.ListColumnsInput`  |
 
 
 ### Example – query a MySQL database
 
-1. Register a connector (either via `dbAddConnection` or through the browser
-   secret-elicitation flow described later). With `dbAddConnection` you pass only
+1. Register a connector (either via `dbSetConnection` or through the browser
+   secret-elicitation flow described later). With `dbSetConnection` you pass only
    **non-secret** fields:
 
 ```jsonc
@@ -317,8 +316,8 @@ sequenceDiagram
     participant UI as "Browser (UI)"
     participant SS as "Secret Store"
 
-    U->>C: dbAddConnection (missing secret)
-    C->>S: rpc dbAddConnection
+    U->>C: dbSetConnection (missing secret)
+    C->>S: rpc dbSetConnection
     S-->>C: Elicit request (flowURI to secret page)
     note over S: entry stored in-memory
     C-->>U: open browser to flowURI
@@ -344,8 +343,8 @@ sequenceDiagram
     participant OP as "OAuth Provider"
     participant SS as "Secret Store"
 
-    U->>C: dbAddConnection (BigQuery etc.)
-    C->>S: rpc dbAddConnection
+    U->>C: dbSetConnection (BigQuery etc.)
+    C->>S: rpc dbSetConnection
     S-->>C: Elicit request (flowURI)
     C-->>U: open browser to flowURI
     U->>UI: navigate

db/connector/add.go

@@ -3,19 +3,21 @@ package connector
 import (
 	"context"
 	"fmt"
-	"github.com/viant/jsonrpc"
 	"net/url"
 	"reflect"
 	"time"
 
+	"github.com/viant/jsonrpc"
+
+	"os"
+	"path/filepath"
+	"strings"
+
 	"github.com/google/uuid"
 	"github.com/viant/mcp-protocol/client"
 	"github.com/viant/mcp-protocol/schema"
 	"github.com/viant/scy"
 	"github.com/viant/scy/cred"
-	"os"
-	"path/filepath"
-	"strings"
 )
 
 // Add registers or updates a connector in the caller's namespace.  If its secret
@@ -24,27 +26,36 @@ import (
 // MCP Elicit protocol, a browser flow is initiated to collect the secret
 // value.  The method never returns the secret and therefore is safe over MCP
 // RPC.
-func (s *Service) Add(ctx context.Context, connector *Connector) error {
+func (s *Service) Add(ctx context.Context, connector *Connector) (*AddOutput, error) {
 	pend, err := s.GeneratePendingSecret(ctx, connector)
 	if err != nil {
-		return err
+		return nil, err
 	}
 	pend.MCP = s.mcpClient
 	connector.secrets = s.secrets
+
 	// If client can handle the Elicit protocol generate it and optionally wait.
 	if impl, ok := s.mcpClient.(client.Operations); ok && impl.Implements(schema.MethodElicitationCreate) {
+		elicitID := uuid.New().String()
+		pend.ElicitID = elicitID
+		oobURL := pend.CallbackURL
+		if strings.Contains(oobURL, "?") {
+			oobURL += "&elicitationId=" + url.QueryEscape(elicitID)
+		} else {
+			oobURL += "?elicitationId=" + url.QueryEscape(elicitID)
+		}
 		elicitResult, _ := impl.Elicit(ctx, &jsonrpc.TypedRequest[*schema.ElicitRequest]{
 			Request: &schema.ElicitRequest{
 				Params: schema.ElicitRequestParams{
-					ElicitationId: uuid.New().String(),
+					ElicitationId: elicitID,
 					Message:       "Open URL to provide secrets for " + connector.Name + " connector",
 					Mode:          "oob",
-					Url:           pend.CallbackURL,
+					Url:           oobURL,
 				}}})
 
 		if elicitResult != nil {
 			if elicitResult.Action != schema.ElicitResultActionAccept {
-				return fmt.Errorf("user reject providing credentials %v", err)
+				return nil, fmt.Errorf("user reject providing credentials %v", err)
 			}
 		}
 		// Wait for secret submission up to 5 min.
@@ -60,10 +71,17 @@ func (s *Service) Add(ctx context.Context, connector *Connector) error {
 				pend.NS.Connectors.Put(connector.Name, connector)
 			}
 		case <-time.After(5 * time.Minute):
+			// Timed out – return pending state with callback URL
+			return &AddOutput{Status: "ok", State: "PENDING_SECRET", CallbackURL: pend.CallbackURL, Connector: connector.Name}, nil
 		case <-ctx.Done():
+			return nil, ctx.Err()
 		}
+		// Secret submitted within wait window – connector activated
+		return &AddOutput{Status: "ok", Connector: connector.Name}, nil
 	}
-	return nil
+	// Client cannot handle Elicit – the connector remains pending waiting for
+	// secret to be supplied out-of-band; return pending state with callback URL.
+	return &AddOutput{Status: "ok", State: "PENDING_SECRET", CallbackURL: pend.CallbackURL, Connector: connector.Name}, nil
 }
 
 func (s *Service) GeneratePendingSecret(ctx context.Context, connector *Connector) (*PendingSecret, error) {

db/connector/connector.go

@@ -3,14 +3,15 @@ package connector
 import (
 	"context"
 	"database/sql"
-	"github.com/viant/mcp-protocol/syncmap"
-	"github.com/viant/scy"
-	"github.com/viant/scy/cred"
 	"reflect"
 	"strings"
 	"sync"
 	"sync/atomic"
 	"time"
+
+	"github.com/viant/mcp-protocol/syncmap"
+	"github.com/viant/scy"
+	"github.com/viant/scy/cred"
 )
 
 type Connectors struct {
@@ -25,7 +26,7 @@ type Connector struct {
 	Name        string        `json:"name" yaml:"name"`
 	DSN         string        `json:"dsn" yaml:"dsn"`
 	Driver      string        `json:"driver" yaml:"driver"`
-	Secrets     *scy.Resource `json:"secrets,omitempty" yaml:"secrets,omitempty"`
+	Secrets     *scy.Resource `json:"secrets,omitempty" yaml:"secrets,omitempty" internal:"true"`
 	db          *sql.DB       `internal:"true"`
 	mux         sync.RWMutex  `internal:"true"`
 	initialized uint32        `internal:"true"`
@@ -55,6 +56,12 @@ func (c *Connector) ExpandDSN(ctx context.Context) (string, error) {
 }
 
 func (c *Connector) Close() error {
+	c.mux.Lock()
+	defer c.mux.Unlock()
+	if c.db != nil {
+		_ = c.db.Close()
+		c.db = nil
+	}
 	atomic.StoreUint32(&c.initialized, 0)
 	return nil
 }

db/connector/elicitation.go

@@ -138,5 +138,8 @@ func (s *Service) requestConnectorElicit(ctx context.Context, impl client.Operat
 		Driver: metaInput.Driver,
 		DSN:    metaInput.Expand(metaConfig.DSN),
 	}
-	return conn.Name, s.Add(ctx, conn)
+	if _, err := s.Add(ctx, conn); err != nil {
+		return "", err
+	}
+	return conn.Name, nil
 }

db/connector/pending.go

@@ -23,6 +23,7 @@ type PendingSecret struct {
 	CredType      reflect.Type
 	MCP           client.Operations
 	OAuth2Config  *oauth2.Config
+	ElicitID      string
 	done          chan struct{}
 }