mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
feat: add aibridgedserver pkg (#19902)
This commit is contained in:
Generated
+2
-1
@@ -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;
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user