688 lines · 24.1 KB · Go
RawDownload
package pluginhost import (	"bytes"	"context"	"encoding/json"	"fmt"	"io"	"net/http"	"strings" 	"github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginabi"	"github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi"	log "github.com/sirupsen/logrus") type rpcPluginAdapter struct {	id     string	host   *Host	client pluginClient} type rpcAuthProvider struct {	*rpcPluginAdapter	identifier string} type rpcFrontendAuthProvider struct {	*rpcPluginAdapter} type rpcProviderExecutor struct {	*rpcPluginAdapter} type rpcThinkingApplier struct {	*rpcPluginAdapter} type rpcQuotaProvider struct {	*rpcPluginAdapter	identifier string} type rpcError struct {	Code       string	message    string	statusCode int} func (e rpcError) Error() string {	return e.message} func (e rpcError) StatusCode() int {	return e.statusCode} type rpcResponseNormalizer struct {	*rpcPluginAdapter	method string} func registerRPCPlugin(ctx context.Context, host *Host, id string, client pluginClient, method string, configYAML []byte) (pluginapi.Plugin, error) {	if client == nil {		return pluginapi.Plugin{}, fmt.Errorf("plugin client is nil")	}	resp, errCall := callPlugin[rpcRegistration](ctx, client, method, rpcLifecycleRequest{		ConfigYAML:    bytes.Clone(configYAML),		SchemaVersion: pluginabi.SchemaVersion,	})	if errCall != nil {		return pluginapi.Plugin{}, errCall	}	if resp.SchemaVersion > pluginabi.SchemaVersion {		return pluginapi.Plugin{}, fmt.Errorf("plugin schema version %d is not supported", resp.SchemaVersion)	}	adapter := &rpcPluginAdapter{id: id, host: host, client: client}	schemaVersion := resp.SchemaVersion	if schemaVersion == 0 {		// Missing schema_version is treated as the original contract.		schemaVersion = 1	}	plugin := pluginapi.Plugin{		Metadata:      resp.Metadata,		SchemaVersion: schemaVersion,		Capabilities: pluginapi.Capabilities{			FrontendAuthProviderExclusive: resp.Capabilities.FrontendAuthProvider && resp.Capabilities.FrontendAuthProviderExclusive,			ExecutorModelScope:            resp.Capabilities.ExecutorModelScope,			ExecutorInputFormats:          append([]string(nil), resp.Capabilities.ExecutorInputFormats...),			ExecutorOutputFormats:         append([]string(nil), resp.Capabilities.ExecutorOutputFormats...),		},	}	if resp.Capabilities.ModelRegistrar {		plugin.Capabilities.ModelRegistrar = adapter	}	if resp.Capabilities.ModelProvider {		plugin.Capabilities.ModelProvider = adapter	}	if resp.Capabilities.AuthProvider {		if err := ctx.Err(); err != nil {			return pluginapi.Plugin{}, err		}		identifier := callPluginIdentifier(ctx, adapter.client, pluginabi.MethodAuthIdentifier)		if err := ctx.Err(); err != nil {			return pluginapi.Plugin{}, err		}		plugin.Capabilities.AuthProvider = rpcAuthProvider{			rpcPluginAdapter: adapter,			identifier:       identifier,		}	}	if resp.Capabilities.FrontendAuthProvider {		plugin.Capabilities.FrontendAuthProvider = rpcFrontendAuthProvider{rpcPluginAdapter: adapter}	}	if resp.Capabilities.Scheduler {		plugin.Capabilities.Scheduler = adapter	}	if resp.Capabilities.ModelRouter {		plugin.Capabilities.ModelRouter = adapter	}	if resp.Capabilities.Executor {		plugin.Capabilities.Executor = rpcProviderExecutor{rpcPluginAdapter: adapter}	}	if resp.Capabilities.RequestTranslator {		plugin.Capabilities.RequestTranslator = adapter	}	if resp.Capabilities.RequestNormalizer {		plugin.Capabilities.RequestNormalizer = adapter	}	if resp.Capabilities.RequestInterceptor {		plugin.Capabilities.RequestInterceptor = adapter	}	if resp.Capabilities.RequestLifecyclePlugin {		plugin.Capabilities.RequestLifecyclePlugin = adapter	}	if resp.Capabilities.ResponseTranslator {		plugin.Capabilities.ResponseTranslator = adapter	}	if resp.Capabilities.ResponseBeforeTranslator {		plugin.Capabilities.ResponseBeforeTranslator = rpcResponseNormalizer{rpcPluginAdapter: adapter, method: pluginabi.MethodResponseNormalizeBefore}	}	if resp.Capabilities.ResponseAfterTranslator {		plugin.Capabilities.ResponseAfterTranslator = rpcResponseNormalizer{rpcPluginAdapter: adapter, method: pluginabi.MethodResponseNormalizeAfter}	}	if resp.Capabilities.ResponseInterceptor {		plugin.Capabilities.ResponseInterceptor = adapter	}	if resp.Capabilities.StreamChunkInterceptor {		plugin.Capabilities.StreamChunkInterceptor = adapter	}	if resp.Capabilities.WebSocketResponseObserver {		plugin.Capabilities.WebSocketResponseObserver = adapter	}	if resp.Capabilities.ThinkingApplier {		plugin.Capabilities.ThinkingApplier = rpcThinkingApplier{rpcPluginAdapter: adapter}	}	if resp.Capabilities.UsagePlugin {		plugin.Capabilities.UsagePlugin = adapter	}	if resp.Capabilities.CommandLinePlugin {		plugin.Capabilities.CommandLinePlugin = adapter	}	if resp.Capabilities.ManagementAPI {		plugin.Capabilities.ManagementAPI = adapter	}	if resp.Capabilities.QuotaProvider {		if err := ctx.Err(); err != nil {			return pluginapi.Plugin{}, err		}		identifier := callPluginIdentifier(ctx, adapter.client, pluginabi.MethodQuotaIdentifier)		if err := ctx.Err(); err != nil {			return pluginapi.Plugin{}, err		}		plugin.Capabilities.QuotaProvider = rpcQuotaProvider{			rpcPluginAdapter: adapter,			identifier:       identifier,		}	}	return plugin, nil} func callPlugin[T any](ctx context.Context, client pluginClient, method string, request any) (T, error) {	var zero T	rawRequest, errMarshal := json.Marshal(sanitizePluginRequest(request))	if errMarshal != nil {		return zero, fmt.Errorf("marshal plugin request %s: %w", method, errMarshal)	}	rawResp, errCall := client.Call(ctx, method, rawRequest)	if errCall != nil {		return zero, errCall	}	var envelope pluginabi.Envelope	if errUnmarshal := json.Unmarshal(rawResp, &envelope); errUnmarshal != nil {		return zero, fmt.Errorf("decode plugin envelope %s: %w", method, errUnmarshal)	}	out, errDecode := decodeEnvelopeResult[T](envelope)	if errDecode != nil {		if !envelope.OK {			return zero, errDecode		}		return zero, fmt.Errorf("decode plugin result %s: %w", method, errDecode)	}	return out, nil} func sanitizePluginRequest(request any) any {	switch req := request.(type) {	case pluginapi.AuthLoginStartRequest:		req.HTTPClient = nil		return req	case pluginapi.AuthLoginPollRequest:		req.HTTPClient = nil		return req	case pluginapi.AuthRefreshRequest:		req.HTTPClient = nil		return req	case pluginapi.AuthModelRequest:		req.HTTPClient = nil		return req	case pluginapi.SchedulerPickRequest:		req.Options.Metadata = sanitizePluginMetadata(req.Options.Metadata)		for index := range req.Candidates {			req.Candidates[index].Metadata = sanitizePluginMetadata(req.Candidates[index].Metadata)		}		return req	case pluginapi.ModelRouteRequest:		req.Metadata = sanitizePluginMetadata(req.Metadata)		return req	case pluginapi.ExecutorRequest:		req.HTTPClient = nil		req.Metadata = sanitizePluginMetadata(req.Metadata)		return req	case pluginapi.RequestInterceptRequest:		req.Metadata = sanitizePluginMetadata(req.Metadata)		return req	case pluginapi.RequestCompletion:		req.Metadata = sanitizePluginMetadata(req.Metadata)		return req	case pluginapi.ResponseInterceptRequest:		req.Metadata = sanitizePluginMetadata(req.Metadata)		return req	case pluginapi.StreamChunkInterceptRequest:		req.Metadata = sanitizePluginMetadata(req.Metadata)		return req	case pluginapi.WebSocketResponseEvent:		req.Metadata = sanitizePluginMetadata(req.Metadata)		return req	case rpcRequestInterceptRequest:		req.Metadata = sanitizePluginMetadata(req.Metadata)		return req	case rpcModelRouteRequest:		req.Metadata = sanitizePluginMetadata(req.Metadata)		return req	case rpcRequestCompletion:		req.Metadata = sanitizePluginMetadata(req.Metadata)		return req	case rpcResponseInterceptRequest:		req.Metadata = sanitizePluginMetadata(req.Metadata)		return req	case rpcStreamChunkInterceptRequest:		req.Metadata = sanitizePluginMetadata(req.Metadata)		return req	case rpcWebSocketResponseEvent:		req.Metadata = sanitizePluginMetadata(req.Metadata)		return req	case pluginapi.ExecutorHTTPRequest:		req.HTTPClient = nil		return req	case rpcExecutorRequest:		req.HTTPClient = nil		req.Metadata = sanitizePluginMetadata(req.Metadata)		return req	case pluginapi.QuotaFetchRequest:		req.HTTPClient = nil		req.Metadata = sanitizePluginMetadata(req.Metadata)		return req	case rpcQuotaFetchRequest:		req.HTTPClient = nil		req.Metadata = sanitizePluginMetadata(req.Metadata)		return req	case pluginapi.QuotaResetRequest:		req.HTTPClient = nil		req.Metadata = sanitizePluginMetadata(req.Metadata)		return req	case rpcQuotaResetRequest:		req.HTTPClient = nil		req.Metadata = sanitizePluginMetadata(req.Metadata)		return req	default:		return request	}} func sanitizePluginMetadata(src map[string]any) map[string]any {	if len(src) == 0 {		return nil	}	dst := make(map[string]any, len(src))	for key, value := range src {		if sanitized, ok := sanitizePluginMetadataValue(value); ok {			dst[key] = sanitized		}	}	if len(dst) == 0 {		return nil	}	return dst} func sanitizePluginMetadataValue(value any) (any, bool) {	switch v := value.(type) {	case nil, string, bool, float64, float32,		int, int8, int16, int32, int64,		uint, uint8, uint16, uint32, uint64:		return value, true	case map[string]any:		return sanitizePluginMetadata(v), true	case []any:		out := make([]any, 0, len(v))		for _, item := range v {			if sanitized, ok := sanitizePluginMetadataValue(item); ok {				out = append(out, sanitized)			}		}		return out, true	default:		// RPC metadata crosses a JSON envelope, so unsupported Go values are normalized to JSON-compatible shapes.		raw, errMarshal := json.Marshal(value)		if errMarshal != nil {			return nil, false		}		var decoded any		if errUnmarshal := json.Unmarshal(raw, &decoded); errUnmarshal != nil {			return nil, false		}		return decoded, true	}} func decodeRPCEnvelope[T any](raw []byte) (T, error) {	var zero T	var envelope pluginabi.Envelope	if errUnmarshal := json.Unmarshal(raw, &envelope); errUnmarshal != nil {		return zero, errUnmarshal	}	return decodeEnvelopeResult[T](envelope)} func isPluginErrorEnvelope(raw []byte) bool {	var envelope pluginabi.Envelope	if errUnmarshal := json.Unmarshal(raw, &envelope); errUnmarshal != nil {		return false	}	return !envelope.OK && envelope.Error != nil} func decodeEnvelopeResult[T any](envelope pluginabi.Envelope) (T, error) {	var zero T	if !envelope.OK {		if envelope.Error != nil {			message := strings.TrimSpace(envelope.Error.Message)			if message == "" {				message = "plugin call failed"			}			return zero, rpcError{				Code:       strings.TrimSpace(envelope.Error.Code),				message:    message,				statusCode: envelope.Error.HTTPStatus,			}		}		return zero, fmt.Errorf("plugin call failed")	}	if len(envelope.Result) == 0 {		return zero, nil	}	var out T	if errDecode := json.Unmarshal(envelope.Result, &out); errDecode != nil {		return zero, errDecode	}	return out, nil} func marshalRPCEnvelope(result json.RawMessage) ([]byte, error) {	if result == nil {		result = json.RawMessage(`{}`)	}	return json.Marshal(pluginabi.Envelope{OK: true, Result: result})} func marshalRPCError(code, message string) []byte {	raw, _ := json.Marshal(pluginabi.Envelope{		OK: false,		Error: &pluginabi.Error{			Code:    code,			Message: message,		},	})	return raw} func (a *rpcPluginAdapter) openHostCallbackContext(ctx context.Context) (string, func()) {	if a == nil || a.host == nil {		return "", func() {}	}	return a.host.openCallbackContextForPlugin(ctx, a.id)} func (a *rpcPluginAdapter) RegisterModels(ctx context.Context, req pluginapi.ModelRegistrationRequest) (pluginapi.ModelRegistrationResponse, error) {	return callPlugin[pluginapi.ModelRegistrationResponse](ctx, a.client, pluginabi.MethodModelRegister, req)} func (a *rpcPluginAdapter) StaticModels(ctx context.Context, req pluginapi.StaticModelRequest) (pluginapi.ModelResponse, error) {	return callPlugin[pluginapi.ModelResponse](ctx, a.client, pluginabi.MethodModelStatic, req)} func (a *rpcPluginAdapter) ModelsForAuth(ctx context.Context, req pluginapi.AuthModelRequest) (pluginapi.ModelResponse, error) {	callbackID, closeCallback := a.openHostCallbackContext(ctx)	defer closeCallback()	return callPlugin[pluginapi.ModelResponse](ctx, a.client, pluginabi.MethodModelForAuth, rpcAuthModelRequest{		AuthModelRequest: req,		HostCallbackID:   callbackID,	})} func (a *rpcPluginAdapter) Pick(ctx context.Context, req pluginapi.SchedulerPickRequest) (pluginapi.SchedulerPickResponse, error) {	return callPlugin[pluginapi.SchedulerPickResponse](ctx, a.client, pluginabi.MethodSchedulerPick, req)} func (a *rpcPluginAdapter) RouteModel(ctx context.Context, req pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, error) {	callbackID, closeCallback := a.openHostCallbackContext(ctx)	defer closeCallback()	return callPlugin[pluginapi.ModelRouteResponse](ctx, a.client, pluginabi.MethodModelRoute, rpcModelRouteRequest{		ModelRouteRequest: req,		HostCallbackID:    callbackID,	})} func callPluginIdentifier(ctx context.Context, client pluginClient, method string) string {	if ctx == nil {		ctx = context.Background()	}	resp, errCall := callPlugin[rpcIdentifierResponse](ctx, client, method, rpcEmptyResponse{})	if errCall != nil {		return ""	}	return strings.TrimSpace(resp.Identifier)} func (a rpcAuthProvider) Identifier() string {	return a.identifier} func (a rpcFrontendAuthProvider) Identifier() string {	return callPluginIdentifier(context.Background(), a.client, pluginabi.MethodFrontendAuthIdentifier)} func (a rpcProviderExecutor) Identifier() string {	return callPluginIdentifier(context.Background(), a.client, pluginabi.MethodExecutorIdentifier)} func (a rpcThinkingApplier) Identifier() string {	return callPluginIdentifier(context.Background(), a.client, pluginabi.MethodThinkingIdentifier)} func (a *rpcPluginAdapter) ParseAuth(ctx context.Context, req pluginapi.AuthParseRequest) (pluginapi.AuthParseResponse, error) {	return callPlugin[pluginapi.AuthParseResponse](ctx, a.client, pluginabi.MethodAuthParse, req)} func (a *rpcPluginAdapter) StartLogin(ctx context.Context, req pluginapi.AuthLoginStartRequest) (pluginapi.AuthLoginStartResponse, error) {	callbackID, closeCallback := a.openHostCallbackContext(ctx)	defer closeCallback()	return callPlugin[pluginapi.AuthLoginStartResponse](ctx, a.client, pluginabi.MethodAuthLoginStart, rpcAuthLoginStartRequest{		AuthLoginStartRequest: req,		HostCallbackID:        callbackID,	})} func (a *rpcPluginAdapter) PollLogin(ctx context.Context, req pluginapi.AuthLoginPollRequest) (pluginapi.AuthLoginPollResponse, error) {	callbackID, closeCallback := a.openHostCallbackContext(ctx)	defer closeCallback()	return callPlugin[pluginapi.AuthLoginPollResponse](ctx, a.client, pluginabi.MethodAuthLoginPoll, rpcAuthLoginPollRequest{		AuthLoginPollRequest: req,		HostCallbackID:       callbackID,	})} func (a *rpcPluginAdapter) RefreshAuth(ctx context.Context, req pluginapi.AuthRefreshRequest) (pluginapi.AuthRefreshResponse, error) {	callbackID, closeCallback := a.openHostCallbackContext(ctx)	defer closeCallback()	return callPlugin[pluginapi.AuthRefreshResponse](ctx, a.client, pluginabi.MethodAuthRefresh, rpcAuthRefreshRequest{		AuthRefreshRequest: req,		HostCallbackID:     callbackID,	})} func (a *rpcPluginAdapter) Authenticate(ctx context.Context, req pluginapi.FrontendAuthRequest) (pluginapi.FrontendAuthResponse, error) {	return callPlugin[pluginapi.FrontendAuthResponse](ctx, a.client, pluginabi.MethodFrontendAuthAuthenticate, req)} 

Lines 1–500 of 688.