Tyk AI Studio Plugin SDK

import "github.com/TykTechnologies/midsommar/v2/pkg/plugin_sdk"

type MyPlugin struct {
    plugin_sdk.BasePlugin
    // Plugin-specific fields
}

func NewMyPlugin() *MyPlugin {
    return &MyPlugin{
        BasePlugin: plugin_sdk.NewBasePlugin(
            "my-plugin",
            "1.0.0",
            "My plugin description",
        ),
    }
}

func (p *MyPlugin) Initialize(ctx plugin_sdk.Context, config map[string]string) error {
    // Initialize plugin
    return nil
}

func main() {
    plugin_sdk.Serve(NewMyPlugin())
}

type RateLimiter struct {
    plugin_sdk.BasePlugin
}

// Implement PostAuthHandler
func (p *RateLimiter) HandlePostAuth(ctx plugin_sdk.Context, req *pb.EnrichedRequest) (*pb.PluginResponse, error) {
    // Check rate limits
}

// Implement ResponseHandler
func (p *RateLimiter) OnBeforeWriteHeaders(ctx plugin_sdk.Context, req *pb.ResponseWriteRequest) (*pb.ResponseWriteResponse, error) {
    // Update counters
}

// Implement UIProvider
func (p *RateLimiter) GetAsset(path string) ([]byte, string, error) {
    // Serve UI assets
}

type MyAuthPlugin struct {
    plugin_sdk.BasePlugin
    tokenStore map[string]*TokenConfig // Maps tokens to app/user IDs
}

// HandleAuth validates the credential and returns App/User IDs
func (p *MyAuthPlugin) HandleAuth(ctx plugin_sdk.Context, req *pb.AuthRequest) (*pb.AuthResponse, error) {
    // Extract token from request
    token := req.Credential
    if token == "" {
        return &pb.AuthResponse{
            Authenticated: false,
            ErrorMessage:  "No credential provided",
        }, nil
    }

    // Validate token and look up associated IDs
    tokenConfig, valid := p.tokenStore[token]
    if !valid {
        return &pb.AuthResponse{
            Authenticated: false,
            ErrorMessage:  "Invalid token",
        }, nil
    }

    // CRITICAL: Return the App ID - this links the credential to access control
    return &pb.AuthResponse{
        Authenticated: true,
        AppId:         tokenConfig.AppID,   // Must be a valid App in the database
        UserId:        tokenConfig.UserID,
    }, nil
}

// GetAppByCredential fetches the App object for access control enforcement
func (p *MyAuthPlugin) GetAppByCredential(ctx plugin_sdk.Context, credential string) (*pb.App, error) {
    tokenConfig, ok := p.tokenStore[credential]
    if !ok {
        return nil, fmt.Errorf("unknown credential")
    }

    // Fetch the App from the Service API
    if ctx.Runtime == plugin_sdk.RuntimeGateway {
        return ctx.Services.Gateway().GetApp(ctx, tokenConfig.AppID)
    }
    return ctx.Services.Studio().GetApp(ctx, tokenConfig.AppID)
}

// GetUserByCredential fetches the User object for identity context
func (p *MyAuthPlugin) GetUserByCredential(ctx plugin_sdk.Context, credential string) (*pb.User, error) {
    tokenConfig, ok := p.tokenStore[credential]
    if !ok {
        return nil, fmt.Errorf("unknown credential")
    }

    // Fetch the User from the Service API
    if ctx.Runtime == plugin_sdk.RuntimeGateway {
        return ctx.Services.Gateway().GetUser(ctx, tokenConfig.UserID)
    }
    return ctx.Services.Studio().GetUser(ctx, tokenConfig.UserID)
}

func (p *MyPlugin) HandlePostAuth(ctx plugin_sdk.Context, req *pb.EnrichedRequest) (*pb.PluginResponse, error) {
    ctx.Services.Logger().Info("Processing request",
        "app_id", ctx.AppID,
        "user_id", ctx.UserID,
    )

    // Add custom header
    req.Headers["X-Custom-Header"] = "value"

    return &pb.PluginResponse{
        Modified: true,
        Request:  req,
    }, nil
}

func (p *MyPlugin) OnBeforeWriteHeaders(ctx plugin_sdk.Context, req *pb.ResponseWriteRequest) (*pb.ResponseWriteResponse, error) {
    // Add tracking header
    if req.Headers == nil {
        req.Headers = make(map[string]string)
    }
    req.Headers["X-Request-Id"] = generateRequestID()

    return &pb.ResponseWriteResponse{
        Modified: true,
        Headers:  req.Headers,
    }, nil
}

func (p *MyPlugin) OnBeforeWrite(ctx plugin_sdk.Context, req *pb.ResponseWriteRequest) (*pb.ResponseWriteResponse, error) {
    // Modify response body
    modifiedBody := transformResponse(req.Body)

    return &pb.ResponseWriteResponse{
        Modified: true,
        Body:     modifiedBody,
    }, nil
}

func (p *MyPlugin) HandleAnalytics(ctx plugin_sdk.Context, analytics *pb.AnalyticsData) error {
    // Export analytics to external system
    return exportToElasticsearch(analytics)
}

func (p *MyPlugin) HandleBudgetUsage(ctx plugin_sdk.Context, usage *pb.BudgetUsageData) error {
    // Track budget usage
    return trackBudget(usage)
}

func (p *MyPlugin) HandleProxyLog(ctx plugin_sdk.Context, log *pb.ProxyLogData) error {
    // Log proxy requests
    return logToFile(log)
}

func (p *MyPlugin) GetConfigSchema() ([]byte, error) {
    schema := map[string]interface{}{
        "type": "object",
        "properties": map[string]interface{}{
            "api_key": map[string]interface{}{
                "type":        "string",
                "description": "API key for external service",
            },
            "rate_limit": map[string]interface{}{
                "type":        "integer",
                "description": "Requests per minute",
                "default":     100,
            },
        },
        "required": []string{"api_key"},
    }
    return json.Marshal(schema)
}

func (p *MyPlugin) AcceptEdgePayload(ctx plugin_sdk.Context, payload *plugin_sdk.EdgePayload) (bool, error) {
    // Check if this payload is for us
    if payload.Metadata["type"] != "my-plugin-data" {
        return false, nil // Not our payload
    }

    ctx.Services.Logger().Info("Received edge payload",
        "edge_id", payload.EdgeID,
        "correlation_id", payload.CorrelationID,
    )

    // Process the payload
    if err := p.processEdgeData(payload.Payload); err != nil {
        return true, err
    }

    return true, nil
}

// Publish an event (flows up from edge to control)
err := ctx.Services.Events().Publish(ctx, "cache.invalidate", payload, plugin_sdk.DirUp)

// Subscribe to events on a specific topic
subscriptionID, err := ctx.Services.Events().Subscribe("cache.invalidate", func(ev plugin_sdk.Event) {
    // Handle event
})

// Subscribe to all events
subscriptionID, err := ctx.Services.Events().SubscribeAll(func(ev plugin_sdk.Event) {
    // Handle any event
})

// Unsubscribe when done
err := ctx.Services.Events().Unsubscribe(subscriptionID)

if ctx.Runtime == plugin_sdk.RuntimeGateway {
    // Get app
    app, err := ctx.Services.Gateway().GetApp(ctx, appID)

    // List apps
    apps, err := ctx.Services.Gateway().ListApps(ctx)

    // Get LLM
    llm, err := ctx.Services.Gateway().GetLLM(ctx, llmID)

    // Get budget status
    status, err := ctx.Services.Gateway().GetBudgetStatus(ctx, appID)

    // Validate credential
    valid, err := ctx.Services.Gateway().ValidateCredential(ctx, token)
}

if ctx.Runtime == plugin_sdk.RuntimeStudio {
    // Get app
    app, err := ctx.Services.Studio().GetApp(ctx, appID)

    // List apps (basic)
    apps, err := ctx.Services.Studio().ListApps(ctx, page, limit)

    // List apps with filtering (by owner, namespace, active status)
    apps, err := ctx.Services.Studio().ListAppsWithFilters(ctx, page, limit, &plugin_sdk.ListAppsOptions{
        UserID: &ownerID,
        Namespace: "production",
    })

    // Update app with metadata (full replacement)
    err := ctx.Services.Studio().UpdateAppWithMetadata(ctx, appID, metadata)

    // Patch a single metadata key (atomic, safe for concurrent use)
    metadataJSON, err := ctx.Services.Studio().PatchAppMetadata(ctx, appID, "tier", `"premium"`, false)

    // List LLMs
    llms, err := ctx.Services.Studio().ListLLMs(ctx, page, limit)

    // List tools
    tools, err := ctx.Services.Studio().ListTools(ctx, page, limit)

    // Call LLM
    stream, err := ctx.Services.Studio().CallLLM(ctx, llmID, model, messages, temp, maxTokens, tools, stream)
}

func (p *MyPlugin) Initialize(ctx plugin_sdk.Context, config map[string]string) error {
    // Extract broker ID for Service API access
    brokerIDStr := ""
    if id, ok := config["_service_broker_id"]; ok {
        brokerIDStr = id
    } else if id, ok := config["service_broker_id"]; ok {
        brokerIDStr = id
    }

    if brokerIDStr != "" {
        var brokerID uint32
        fmt.Sscanf(brokerIDStr, "%d", &brokerID)
        ai_studio_sdk.SetServiceBrokerID(brokerID)
    }

    // Parse plugin-specific config
    p.apiKey = config["api_key"]

    return nil
}

type MyPlugin struct {
    plugin_sdk.BasePlugin
    apiKey string
}

func NewMyPlugin() *MyPlugin {
    return &MyPlugin{
        BasePlugin: plugin_sdk.NewBasePlugin(
            "my-plugin",
            "1.0.0",
            "My plugin description",
        ),
    }
}

package main

import (
    "encoding/json"
    "github.com/TykTechnologies/midsommar/v2/pkg/plugin_sdk"
    "github.com/TykTechnologies/midsommar/v2/pkg/plugin_sdk/pb"
)

type RequestLogger struct {
    plugin_sdk.BasePlugin
}

func NewRequestLogger() *RequestLogger {
    return &RequestLogger{
        BasePlugin: plugin_sdk.NewBasePlugin(
            "request-logger",
            "1.0.0",
            "Logs requests and responses",
        ),
    }
}

// PostAuthHandler: Log incoming requests
func (p *RequestLogger) HandlePostAuth(ctx plugin_sdk.Context, req *pb.EnrichedRequest) (*pb.PluginResponse, error) {
    ctx.Services.Logger().Info("Incoming request",
        "app_id", ctx.AppID,
        "user_id", ctx.UserID,
        "path", req.Path,
        "method", req.Method,
    )

    // Store request metadata in KV
    metadata := map[string]interface{}{
        "timestamp": time.Now().Unix(),
        "path":      req.Path,
        "method":    req.Method,
    }
    data, _ := json.Marshal(metadata)
    ctx.Services.KV().Write(ctx, fmt.Sprintf("req:%s", req.RequestId), data)

    return &pb.PluginResponse{Modified: false}, nil
}

// ResponseHandler: Log responses
func (p *RequestLogger) OnBeforeWriteHeaders(ctx plugin_sdk.Context, req *pb.ResponseWriteRequest) (*pb.ResponseWriteResponse, error) {
    ctx.Services.Logger().Info("Outgoing response",
        "app_id", ctx.AppID,
        "status_code", req.StatusCode,
        "request_id", req.RequestId,
    )

    return &pb.ResponseWriteResponse{Modified: false}, nil
}

// ConfigProvider: Provide config schema
func (p *RequestLogger) GetConfigSchema() ([]byte, error) {
    schema := map[string]interface{}{
        "type": "object",
        "properties": map[string]interface{}{
            "log_level": map[string]interface{}{
                "type":    "string",
                "enum":    []string{"debug", "info", "warn", "error"},
                "default": "info",
            },
        },
    }
    return json.Marshal(schema)
}

func main() {
    plugin_sdk.Serve(NewRequestLogger())
}

func TestPluginLogic(t *testing.T) {
    plugin := NewRequestLogger()

    ctx := plugin_sdk.Context{
        Runtime: plugin_sdk.RuntimeStudio,
        AppID:   1,
    }

    req := &pb.EnrichedRequest{
        Path:   "/api/v1/chat",
        Method: "POST",
    }

    resp, err := plugin.HandlePostAuth(ctx, req)
    if err != nil {
        t.Fatalf("HandlePostAuth failed: %v", err)
    }

    if resp.Block {
        t.Error("Expected request to not be blocked")
    }
}

type MyPlugin struct {
    plugin_sdk.BasePlugin
    services plugin_sdk.ServiceBroker
}

func (p *MyPlugin) Initialize(ctx plugin_sdk.Context, config map[string]string) error {
    p.services = ctx.Services
    return nil
}

// OnSessionReady implements plugin_sdk.SessionAware
// This is called when the session-based broker connection is established.
func (p *MyPlugin) OnSessionReady(ctx plugin_sdk.Context) {
    log.Printf("Session ready - warming up service API connection...")

    // Eagerly establish the broker connection by making a simple API call.
    // This "warms up" the connection so subsequent RPC calls don't need to dial.
    if ai_studio_sdk.IsInitialized() {
        // Make a lightweight API call to establish the connection
        _, err := ai_studio_sdk.GetPluginsCount(context.Background())
        if err != nil {
            log.Printf("Service API warmup failed: %v", err)
        } else {
            log.Printf("Service API connection established successfully")
        }
    }
}

// OnSessionClosing implements plugin_sdk.SessionAware
func (p *MyPlugin) OnSessionClosing(ctx plugin_sdk.Context) {
    log.Printf("Session closing - cleaning up resources")
}

package main

import (
    "context"
    "log"

    "github.com/TykTechnologies/midsommar/v2/pkg/ai_studio_sdk"
    "github.com/TykTechnologies/midsommar/v2/pkg/plugin_sdk"
)

type MyServicePlugin struct {
    plugin_sdk.BasePlugin
    services      plugin_sdk.ServiceBroker
    eventSubID    string
}

func NewMyServicePlugin() *MyServicePlugin {
    return &MyServicePlugin{
        BasePlugin: plugin_sdk.NewBasePlugin("my-service-plugin", "1.0.0", "Plugin using Service API"),
    }
}

func (p *MyServicePlugin) Initialize(ctx plugin_sdk.Context, config map[string]string) error {
    p.services = ctx.Services
    log.Printf("Initialized in %s runtime", ctx.Runtime)
    return nil
}

// OnSessionReady - warm up connections and set up subscriptions
func (p *MyServicePlugin) OnSessionReady(ctx plugin_sdk.Context) {
    log.Printf("Session ready")

    // 1. Warm up Service API connection
    if ai_studio_sdk.IsInitialized() {
        _, err := ai_studio_sdk.GetPluginsCount(context.Background())
        if err != nil {
            log.Printf("Service API warmup failed: %v", err)
        } else {
            log.Printf("Service API connection ready")
        }
    }

    // 2. Set up event subscriptions (uses same connection)
    if p.services != nil {
        events := p.services.Events()
        if events != nil {
            subID, err := events.Subscribe("config.updated", p.handleConfigUpdate)
            if err != nil {
                log.Printf("Failed to subscribe to events: %v", err)
            } else {
                p.eventSubID = subID
                log.Printf("Subscribed to config.updated events")
            }
        }
    }
}

func (p *MyServicePlugin) handleConfigUpdate(ev plugin_sdk.Event) {
    log.Printf("Received config update: %s", ev.Topic)
    // Handle the event...
}

func (p *MyServicePlugin) OnSessionClosing(ctx plugin_sdk.Context) {
    // Clean up event subscription
    if p.eventSubID != "" && p.services != nil {
        p.services.Events().Unsubscribe(p.eventSubID)
    }
}

// HandleRPC - called from UI, Service API is already warmed up
func (p *MyServicePlugin) HandleRPC(method string, payload []byte) ([]byte, error) {
    // Service API calls will work because connection was warmed up in OnSessionReady
    llms, err := ai_studio_sdk.ListLLMs(context.Background(), 1, 10)
    if err != nil {
        return nil, err
    }
    // Process llms...
    return []byte(`{"success": true}`), nil
}

func main() {
    plugin_sdk.Serve(NewMyServicePlugin())
}