feat: add aibridgedserver pkg (#19902)

This commit is contained in:
Danny Kopping
2025-09-25 13:32:16 +02:00
committed by GitHub
parent c8742badbb
commit 615585d5d1
15 changed files with 1587 additions and 347 deletions
+2 -1
View File
@@ -990,7 +990,8 @@ CREATE TABLE aibridge_interceptions (
initiator_id uuid NOT NULL,
provider text NOT NULL,
model text NOT NULL,
started_at timestamp with time zone NOT NULL
started_at timestamp with time zone NOT NULL,
metadata jsonb
);
COMMENT ON TABLE aibridge_interceptions IS 'Audit log of requests intercepted by AI Bridge';
@@ -0,0 +1 @@
ALTER TABLE aibridge_interceptions DROP COLUMN metadata;
@@ -0,0 +1 @@
ALTER TABLE aibridge_interceptions ADD COLUMN metadata JSONB DEFAULT NULL;
+5 -4
View File
@@ -3373,10 +3373,11 @@ func AllWorkspaceTransitionValues() []WorkspaceTransition {
type AIBridgeInterception struct {
ID uuid.UUID `db:"id" json:"id"`
// Relates to a users record, but FK is elided for performance.
InitiatorID uuid.UUID `db:"initiator_id" json:"initiator_id"`
Provider string `db:"provider" json:"provider"`
Model string `db:"model" json:"model"`
StartedAt time.Time `db:"started_at" json:"started_at"`
InitiatorID uuid.UUID `db:"initiator_id" json:"initiator_id"`
Provider string `db:"provider" json:"provider"`
Model string `db:"model" json:"model"`
StartedAt time.Time `db:"started_at" json:"started_at"`
Metadata pqtype.NullRawMessage `db:"metadata" json:"metadata"`
}
// Audit log of tokens used by intercepted requests in AI Bridge
+13 -9
View File
@@ -112,7 +112,7 @@ func (q *sqlQuerier) ActivityBumpWorkspace(ctx context.Context, arg ActivityBump
}
const getAIBridgeInterceptionByID = `-- name: GetAIBridgeInterceptionByID :one
SELECT id, initiator_id, provider, model, started_at FROM aibridge_interceptions WHERE id = $1::uuid
SELECT id, initiator_id, provider, model, started_at, metadata FROM aibridge_interceptions WHERE id = $1::uuid
`
func (q *sqlQuerier) GetAIBridgeInterceptionByID(ctx context.Context, id uuid.UUID) (AIBridgeInterception, error) {
@@ -124,22 +124,24 @@ func (q *sqlQuerier) GetAIBridgeInterceptionByID(ctx context.Context, id uuid.UU
&i.Provider,
&i.Model,
&i.StartedAt,
&i.Metadata,
)
return i, err
}
const insertAIBridgeInterception = `-- name: InsertAIBridgeInterception :one
INSERT INTO aibridge_interceptions (id, initiator_id, provider, model, started_at)
VALUES ($1::uuid, $2::uuid, $3, $4, $5)
RETURNING id, initiator_id, provider, model, started_at
INSERT INTO aibridge_interceptions (id, initiator_id, provider, model, metadata, started_at)
VALUES ($1::uuid, $2::uuid, $3, $4, COALESCE($5::jsonb, '{}'::jsonb), $6)
RETURNING id, initiator_id, provider, model, started_at, metadata
`
type InsertAIBridgeInterceptionParams struct {
ID uuid.UUID `db:"id" json:"id"`
InitiatorID uuid.UUID `db:"initiator_id" json:"initiator_id"`
Provider string `db:"provider" json:"provider"`
Model string `db:"model" json:"model"`
StartedAt time.Time `db:"started_at" json:"started_at"`
ID uuid.UUID `db:"id" json:"id"`
InitiatorID uuid.UUID `db:"initiator_id" json:"initiator_id"`
Provider string `db:"provider" json:"provider"`
Model string `db:"model" json:"model"`
Metadata json.RawMessage `db:"metadata" json:"metadata"`
StartedAt time.Time `db:"started_at" json:"started_at"`
}
func (q *sqlQuerier) InsertAIBridgeInterception(ctx context.Context, arg InsertAIBridgeInterceptionParams) (AIBridgeInterception, error) {
@@ -148,6 +150,7 @@ func (q *sqlQuerier) InsertAIBridgeInterception(ctx context.Context, arg InsertA
arg.InitiatorID,
arg.Provider,
arg.Model,
arg.Metadata,
arg.StartedAt,
)
var i AIBridgeInterception
@@ -157,6 +160,7 @@ func (q *sqlQuerier) InsertAIBridgeInterception(ctx context.Context, arg InsertA
&i.Provider,
&i.Model,
&i.StartedAt,
&i.Metadata,
)
return i, err
}
+2 -2
View File
@@ -1,6 +1,6 @@
-- name: InsertAIBridgeInterception :one
INSERT INTO aibridge_interceptions (id, initiator_id, provider, model, started_at)
VALUES (@id::uuid, @initiator_id::uuid, @provider, @model, @started_at)
INSERT INTO aibridge_interceptions (id, initiator_id, provider, model, metadata, started_at)
VALUES (@id::uuid, @initiator_id::uuid, @provider, @model, COALESCE(@metadata::jsonb, '{}'::jsonb), @started_at)
RETURNING *;
-- name: InsertAIBridgeTokenUsage :exec
+3
View File
@@ -24,6 +24,9 @@ const (
MCPServerName = "Coder"
// MCPServerInstructions is the instructions text for the MCP server.
MCPServerInstructions = "Coder MCP Server providing workspace and template management tools"
// Used in tests and aibridge.
MCPEndpoint = "/api/experimental/mcp/http"
)
// Server represents an MCP HTTP server instance
+11 -11
View File
@@ -34,7 +34,7 @@ func TestMCPHTTP_E2E_ClientIntegration(t *testing.T) {
_ = coderdtest.CreateFirstUser(t, coderClient)
// Create MCP client pointing to our endpoint
mcpURL := api.AccessURL.String() + "/api/experimental/mcp/http"
mcpURL := api.AccessURL.String() + mcpserver.MCPEndpoint
// Configure client with authentication headers using RFC 6750 Bearer token
mcpClient, err := mcpclient.NewStreamableHttpClient(mcpURL,
@@ -133,7 +133,7 @@ func TestMCPHTTP_E2E_UnauthenticatedAccess(t *testing.T) {
defer cancel()
// Test direct HTTP request to verify 401 status code
mcpURL := api.AccessURL.String() + "/api/experimental/mcp/http"
mcpURL := api.AccessURL.String() + mcpserver.MCPEndpoint
// Make a POST request without authentication (MCP over HTTP uses POST)
//nolint:gosec // Test code using controlled localhost URL
@@ -197,7 +197,7 @@ func TestMCPHTTP_E2E_ToolWithWorkspace(t *testing.T) {
workspace := coderdtest.CreateWorkspace(t, coderClient, template.ID)
// Create MCP client
mcpURL := api.AccessURL.String() + "/api/experimental/mcp/http"
mcpURL := api.AccessURL.String() + mcpserver.MCPEndpoint
mcpClient, err := mcpclient.NewStreamableHttpClient(mcpURL,
transport.WithHTTPHeaders(map[string]string{
"Authorization": "Bearer " + coderClient.SessionToken(),
@@ -280,7 +280,7 @@ func TestMCPHTTP_E2E_ErrorHandling(t *testing.T) {
_ = coderdtest.CreateFirstUser(t, coderClient)
// Create MCP client
mcpURL := api.AccessURL.String() + "/api/experimental/mcp/http"
mcpURL := api.AccessURL.String() + mcpserver.MCPEndpoint
mcpClient, err := mcpclient.NewStreamableHttpClient(mcpURL,
transport.WithHTTPHeaders(map[string]string{
"Authorization": "Bearer " + coderClient.SessionToken(),
@@ -339,7 +339,7 @@ func TestMCPHTTP_E2E_ConcurrentRequests(t *testing.T) {
_ = coderdtest.CreateFirstUser(t, coderClient)
// Create MCP client
mcpURL := api.AccessURL.String() + "/api/experimental/mcp/http"
mcpURL := api.AccessURL.String() + mcpserver.MCPEndpoint
mcpClient, err := mcpclient.NewStreamableHttpClient(mcpURL,
transport.WithHTTPHeaders(map[string]string{
"Authorization": "Bearer " + coderClient.SessionToken(),
@@ -410,7 +410,7 @@ func TestMCPHTTP_E2E_RFC6750_UnauthenticatedRequest(t *testing.T) {
// Make a request without any authentication headers
req := &http.Request{
Method: "POST",
URL: mustParseURL(t, api.AccessURL.String()+"/api/experimental/mcp/http"),
URL: mustParseURL(t, api.AccessURL.String()+mcpserver.MCPEndpoint),
Header: make(http.Header),
}
@@ -493,7 +493,7 @@ func TestMCPHTTP_E2E_OAuth2_EndToEnd(t *testing.T) {
// In a real OAuth2 flow, this would be an OAuth2 access token
sessionToken := coderClient.SessionToken()
mcpURL := api.AccessURL.String() + "/api/experimental/mcp/http"
mcpURL := api.AccessURL.String() + mcpserver.MCPEndpoint
mcpClient, err := mcpclient.NewStreamableHttpClient(mcpURL,
transport.WithHTTPHeaders(map[string]string{
"Authorization": "Bearer " + sessionToken,
@@ -640,7 +640,7 @@ func TestMCPHTTP_E2E_OAuth2_EndToEnd(t *testing.T) {
t.Logf("Successfully obtained refresh token: %s...", refreshToken[:10])
// Step 3: Use access token to authenticate with MCP endpoint
mcpURL := api.AccessURL.String() + "/api/experimental/mcp/http"
mcpURL := api.AccessURL.String() + mcpserver.MCPEndpoint
mcpClient, err := mcpclient.NewStreamableHttpClient(mcpURL,
transport.WithHTTPHeaders(map[string]string{
"Authorization": "Bearer " + accessToken,
@@ -776,7 +776,7 @@ func TestMCPHTTP_E2E_OAuth2_EndToEnd(t *testing.T) {
t.Parallel()
req := &http.Request{
Method: "POST",
URL: mustParseURL(t, api.AccessURL.String()+"/api/experimental/mcp/http"),
URL: mustParseURL(t, api.AccessURL.String()+mcpserver.MCPEndpoint),
Header: map[string][]string{
"Authorization": {"Bearer invalid_token_value"},
"Content-Type": {"application/json"},
@@ -805,7 +805,7 @@ func TestMCPHTTP_E2E_OAuth2_EndToEnd(t *testing.T) {
t.Run("DynamicClientRegistrationWithMCPFlow", func(t *testing.T) {
t.Parallel()
// Step 1: Attempt unauthenticated MCP access
mcpURL := api.AccessURL.String() + "/api/experimental/mcp/http"
mcpURL := api.AccessURL.String() + mcpserver.MCPEndpoint
req := &http.Request{
Method: "POST",
URL: mustParseURL(t, mcpURL),
@@ -1232,7 +1232,7 @@ func TestMCPHTTP_E2E_ChatGPTEndpoint(t *testing.T) {
template := coderdtest.CreateTemplate(t, coderClient, user.OrganizationID, version.ID)
// Create MCP client pointing to the ChatGPT endpoint
mcpURL := api.AccessURL.String() + "/api/experimental/mcp/http?toolset=chatgpt"
mcpURL := api.AccessURL.String() + mcpserver.MCPEndpoint + "?toolset=chatgpt"
// Configure client with authentication headers using RFC 6750 Bearer token
mcpClient, err := mcpclient.NewStreamableHttpClient(mcpURL,