chore: define aibridged protobuf specification (#19795)

This commit is contained in:
Danny Kopping
2025-09-24 17:05:04 +10:00
committed by GitHub
parent 5f1ddb6d57
commit 0ece153a0e
4 changed files with 1967 additions and 0 deletions
+10
View File
@@ -636,6 +636,7 @@ GEN_FILES := \
provisionersdk/proto/provisioner.pb.go \
provisionerd/proto/provisionerd.pb.go \
vpn/vpn.pb.go \
aibridged/proto/aibridged.pb.go \
$(DB_GEN_FILES) \
$(SITE_GEN_FILES) \
coderd/rbac/object_gen.go \
@@ -683,6 +684,7 @@ gen/mark-fresh:
provisionersdk/proto/provisioner.pb.go \
provisionerd/proto/provisionerd.pb.go \
vpn/vpn.pb.go \
aibridged/proto/aibridged.pb.go \
coderd/database/dump.sql \
$(DB_GEN_FILES) \
site/src/api/typesGenerated.ts \
@@ -801,6 +803,14 @@ vpn/vpn.pb.go: vpn/vpn.proto
--go_opt=paths=source_relative \
./vpn/vpn.proto
aibridged/proto/aibridged.pb.go: aibridged/proto/aibridged.proto
protoc \
--go_out=. \
--go_opt=paths=source_relative \
--go-drpc_out=. \
--go-drpc_opt=paths=source_relative \
./aibridged/proto/aibridged.proto
site/src/api/typesGenerated.ts: site/node_modules/.installed $(wildcard scripts/apitypings/*) $(shell find ./codersdk $(FIND_EXCLUSIONS) -type f -name '*.go')
# -C sets the directory for the go run command
go run -C ./scripts/apitypings main.go > $@
File diff suppressed because it is too large Load Diff
+114
View File
@@ -0,0 +1,114 @@
syntax = "proto3";
option go_package = "github.com/coder/coder/v2/aibridged/proto";
package proto;
import "google/protobuf/any.proto";
import "google/protobuf/timestamp.proto";
// Recorder is responsible for persisting AI usage records along with their related interception.
service Recorder {
// RecordInterception creates a new interception record to which all other sub-resources
// (token, prompt, tool uses) will be related.
rpc RecordInterception(RecordInterceptionRequest) returns (RecordInterceptionResponse);
rpc RecordTokenUsage(RecordTokenUsageRequest) returns (RecordTokenUsageResponse);
rpc RecordPromptUsage(RecordPromptUsageRequest) returns (RecordPromptUsageResponse);
rpc RecordToolUsage(RecordToolUsageRequest) returns (RecordToolUsageResponse);
}
// MCPConfigurator is responsible for retrieving any relevant data required for configuring MCP clients
// against remote servers.
service MCPConfigurator {
// GetMCPServerConfigs will retrieve MCP server configurations.
rpc GetMCPServerConfigs(GetMCPServerConfigsRequest) returns (GetMCPServerConfigsResponse);
// GetMCPServerAccessTokensBatch will retrieve an access token for a given list of MCP servers, which may involve
// acquiring, validating, or refreshing tokens synchronously. The server should make every effort to
// parallelise this work.
rpc GetMCPServerAccessTokensBatch(GetMCPServerAccessTokensBatchRequest) returns (GetMCPServerAccessTokensBatchResponse);
}
// Authorizer handles all Coder-related authorization functions.
service Authorizer {
// IsAuthorized validates that a given Coder key is valid and the user is authorized to use AI Bridge.
// TODO: add authorization; currently only key validation takes place.
rpc IsAuthorized(IsAuthorizedRequest) returns (IsAuthorizedResponse);
}
message RecordInterceptionRequest {
string id = 1; // UUID.
string initiator_id = 2; // UUID.
string provider = 3;
string model = 4;
map<string, google.protobuf.Any> metadata = 5;
google.protobuf.Timestamp started_at = 6;
}
message RecordInterceptionResponse {}
message RecordTokenUsageRequest {
string interception_id = 1; // UUID.
string msg_id = 2; // ID provided by provider.
int64 input_tokens = 3;
int64 output_tokens = 4;
map<string, google.protobuf.Any> metadata = 5;
google.protobuf.Timestamp created_at = 6;
}
message RecordTokenUsageResponse {}
message RecordPromptUsageRequest {
string interception_id = 1; // UUID.
string msg_id = 2; // ID provided by provider.
string prompt = 3;
map<string, google.protobuf.Any> metadata = 4;
google.protobuf.Timestamp created_at = 5;
}
message RecordPromptUsageResponse {}
message RecordToolUsageRequest {
string interception_id = 1; // UUID.
string msg_id = 2; // ID provided by provider.
optional string server_url = 3; // The URL of the MCP server.
string tool = 4;
string input = 5;
bool injected = 6;
optional string invocation_error = 7; // Only injected tools are invoked.
map<string, google.protobuf.Any> metadata = 8;
google.protobuf.Timestamp created_at = 9;
}
message RecordToolUsageResponse {}
message GetMCPServerConfigsRequest {
string user_id = 1; // UUID. // Not used yet, will be necessary for later RBAC purposes.
}
message GetMCPServerConfigsResponse {
MCPServerConfig coder_mcp_config = 1;
repeated MCPServerConfig external_auth_mcp_configs = 2;
}
message MCPServerConfig {
string id = 1; // Maps to the ID of the External Auth; this ID is unique.
string url = 2;
string tool_allow_regex = 3;
string tool_deny_regex = 4;
}
message GetMCPServerAccessTokensBatchRequest {
string user_id = 1; // UUID.
repeated string mcp_server_config_ids = 2;
}
// GetMCPServerAccessTokensBatchResponse returns a map for resulting tokens or errors, indexed
// by server ID.
message GetMCPServerAccessTokensBatchResponse{
map<string, string> access_tokens = 1;
map<string, string> errors = 2;
}
message IsAuthorizedRequest {
string key = 1;
}
message IsAuthorizedResponse {
string owner_id = 1;
}
+421
View File
@@ -0,0 +1,421 @@
// Code generated by protoc-gen-go-drpc. DO NOT EDIT.
// protoc-gen-go-drpc version: v0.0.34
// source: aibridged/proto/aibridged.proto
package proto
import (
context "context"
errors "errors"
protojson "google.golang.org/protobuf/encoding/protojson"
proto "google.golang.org/protobuf/proto"
drpc "storj.io/drpc"
drpcerr "storj.io/drpc/drpcerr"
)
type drpcEncoding_File_aibridged_proto_aibridged_proto struct{}
func (drpcEncoding_File_aibridged_proto_aibridged_proto) Marshal(msg drpc.Message) ([]byte, error) {
return proto.Marshal(msg.(proto.Message))
}
func (drpcEncoding_File_aibridged_proto_aibridged_proto) MarshalAppend(buf []byte, msg drpc.Message) ([]byte, error) {
return proto.MarshalOptions{}.MarshalAppend(buf, msg.(proto.Message))
}
func (drpcEncoding_File_aibridged_proto_aibridged_proto) Unmarshal(buf []byte, msg drpc.Message) error {
return proto.Unmarshal(buf, msg.(proto.Message))
}
func (drpcEncoding_File_aibridged_proto_aibridged_proto) JSONMarshal(msg drpc.Message) ([]byte, error) {
return protojson.Marshal(msg.(proto.Message))
}
func (drpcEncoding_File_aibridged_proto_aibridged_proto) JSONUnmarshal(buf []byte, msg drpc.Message) error {
return protojson.Unmarshal(buf, msg.(proto.Message))
}
type DRPCRecorderClient interface {
DRPCConn() drpc.Conn
RecordInterception(ctx context.Context, in *RecordInterceptionRequest) (*RecordInterceptionResponse, error)
RecordTokenUsage(ctx context.Context, in *RecordTokenUsageRequest) (*RecordTokenUsageResponse, error)
RecordPromptUsage(ctx context.Context, in *RecordPromptUsageRequest) (*RecordPromptUsageResponse, error)
RecordToolUsage(ctx context.Context, in *RecordToolUsageRequest) (*RecordToolUsageResponse, error)
}
type drpcRecorderClient struct {
cc drpc.Conn
}
func NewDRPCRecorderClient(cc drpc.Conn) DRPCRecorderClient {
return &drpcRecorderClient{cc}
}
func (c *drpcRecorderClient) DRPCConn() drpc.Conn { return c.cc }
func (c *drpcRecorderClient) RecordInterception(ctx context.Context, in *RecordInterceptionRequest) (*RecordInterceptionResponse, error) {
out := new(RecordInterceptionResponse)
err := c.cc.Invoke(ctx, "/proto.Recorder/RecordInterception", drpcEncoding_File_aibridged_proto_aibridged_proto{}, in, out)
if err != nil {
return nil, err
}
return out, nil
}
func (c *drpcRecorderClient) RecordTokenUsage(ctx context.Context, in *RecordTokenUsageRequest) (*RecordTokenUsageResponse, error) {
out := new(RecordTokenUsageResponse)
err := c.cc.Invoke(ctx, "/proto.Recorder/RecordTokenUsage", drpcEncoding_File_aibridged_proto_aibridged_proto{}, in, out)
if err != nil {
return nil, err
}
return out, nil
}
func (c *drpcRecorderClient) RecordPromptUsage(ctx context.Context, in *RecordPromptUsageRequest) (*RecordPromptUsageResponse, error) {
out := new(RecordPromptUsageResponse)
err := c.cc.Invoke(ctx, "/proto.Recorder/RecordPromptUsage", drpcEncoding_File_aibridged_proto_aibridged_proto{}, in, out)
if err != nil {
return nil, err
}
return out, nil
}
func (c *drpcRecorderClient) RecordToolUsage(ctx context.Context, in *RecordToolUsageRequest) (*RecordToolUsageResponse, error) {
out := new(RecordToolUsageResponse)
err := c.cc.Invoke(ctx, "/proto.Recorder/RecordToolUsage", drpcEncoding_File_aibridged_proto_aibridged_proto{}, in, out)
if err != nil {
return nil, err
}
return out, nil
}
type DRPCRecorderServer interface {
RecordInterception(context.Context, *RecordInterceptionRequest) (*RecordInterceptionResponse, error)
RecordTokenUsage(context.Context, *RecordTokenUsageRequest) (*RecordTokenUsageResponse, error)
RecordPromptUsage(context.Context, *RecordPromptUsageRequest) (*RecordPromptUsageResponse, error)
RecordToolUsage(context.Context, *RecordToolUsageRequest) (*RecordToolUsageResponse, error)
}
type DRPCRecorderUnimplementedServer struct{}
func (s *DRPCRecorderUnimplementedServer) RecordInterception(context.Context, *RecordInterceptionRequest) (*RecordInterceptionResponse, error) {
return nil, drpcerr.WithCode(errors.New("Unimplemented"), drpcerr.Unimplemented)
}
func (s *DRPCRecorderUnimplementedServer) RecordTokenUsage(context.Context, *RecordTokenUsageRequest) (*RecordTokenUsageResponse, error) {
return nil, drpcerr.WithCode(errors.New("Unimplemented"), drpcerr.Unimplemented)
}
func (s *DRPCRecorderUnimplementedServer) RecordPromptUsage(context.Context, *RecordPromptUsageRequest) (*RecordPromptUsageResponse, error) {
return nil, drpcerr.WithCode(errors.New("Unimplemented"), drpcerr.Unimplemented)
}
func (s *DRPCRecorderUnimplementedServer) RecordToolUsage(context.Context, *RecordToolUsageRequest) (*RecordToolUsageResponse, error) {
return nil, drpcerr.WithCode(errors.New("Unimplemented"), drpcerr.Unimplemented)
}
type DRPCRecorderDescription struct{}
func (DRPCRecorderDescription) NumMethods() int { return 4 }
func (DRPCRecorderDescription) Method(n int) (string, drpc.Encoding, drpc.Receiver, interface{}, bool) {
switch n {
case 0:
return "/proto.Recorder/RecordInterception", drpcEncoding_File_aibridged_proto_aibridged_proto{},
func(srv interface{}, ctx context.Context, in1, in2 interface{}) (drpc.Message, error) {
return srv.(DRPCRecorderServer).
RecordInterception(
ctx,
in1.(*RecordInterceptionRequest),
)
}, DRPCRecorderServer.RecordInterception, true
case 1:
return "/proto.Recorder/RecordTokenUsage", drpcEncoding_File_aibridged_proto_aibridged_proto{},
func(srv interface{}, ctx context.Context, in1, in2 interface{}) (drpc.Message, error) {
return srv.(DRPCRecorderServer).
RecordTokenUsage(
ctx,
in1.(*RecordTokenUsageRequest),
)
}, DRPCRecorderServer.RecordTokenUsage, true
case 2:
return "/proto.Recorder/RecordPromptUsage", drpcEncoding_File_aibridged_proto_aibridged_proto{},
func(srv interface{}, ctx context.Context, in1, in2 interface{}) (drpc.Message, error) {
return srv.(DRPCRecorderServer).
RecordPromptUsage(
ctx,
in1.(*RecordPromptUsageRequest),
)
}, DRPCRecorderServer.RecordPromptUsage, true
case 3:
return "/proto.Recorder/RecordToolUsage", drpcEncoding_File_aibridged_proto_aibridged_proto{},
func(srv interface{}, ctx context.Context, in1, in2 interface{}) (drpc.Message, error) {
return srv.(DRPCRecorderServer).
RecordToolUsage(
ctx,
in1.(*RecordToolUsageRequest),
)
}, DRPCRecorderServer.RecordToolUsage, true
default:
return "", nil, nil, nil, false
}
}
func DRPCRegisterRecorder(mux drpc.Mux, impl DRPCRecorderServer) error {
return mux.Register(impl, DRPCRecorderDescription{})
}
type DRPCRecorder_RecordInterceptionStream interface {
drpc.Stream
SendAndClose(*RecordInterceptionResponse) error
}
type drpcRecorder_RecordInterceptionStream struct {
drpc.Stream
}
func (x *drpcRecorder_RecordInterceptionStream) SendAndClose(m *RecordInterceptionResponse) error {
if err := x.MsgSend(m, drpcEncoding_File_aibridged_proto_aibridged_proto{}); err != nil {
return err
}
return x.CloseSend()
}
type DRPCRecorder_RecordTokenUsageStream interface {
drpc.Stream
SendAndClose(*RecordTokenUsageResponse) error
}
type drpcRecorder_RecordTokenUsageStream struct {
drpc.Stream
}
func (x *drpcRecorder_RecordTokenUsageStream) SendAndClose(m *RecordTokenUsageResponse) error {
if err := x.MsgSend(m, drpcEncoding_File_aibridged_proto_aibridged_proto{}); err != nil {
return err
}
return x.CloseSend()
}
type DRPCRecorder_RecordPromptUsageStream interface {
drpc.Stream
SendAndClose(*RecordPromptUsageResponse) error
}
type drpcRecorder_RecordPromptUsageStream struct {
drpc.Stream
}
func (x *drpcRecorder_RecordPromptUsageStream) SendAndClose(m *RecordPromptUsageResponse) error {
if err := x.MsgSend(m, drpcEncoding_File_aibridged_proto_aibridged_proto{}); err != nil {
return err
}
return x.CloseSend()
}
type DRPCRecorder_RecordToolUsageStream interface {
drpc.Stream
SendAndClose(*RecordToolUsageResponse) error
}
type drpcRecorder_RecordToolUsageStream struct {
drpc.Stream
}
func (x *drpcRecorder_RecordToolUsageStream) SendAndClose(m *RecordToolUsageResponse) error {
if err := x.MsgSend(m, drpcEncoding_File_aibridged_proto_aibridged_proto{}); err != nil {
return err
}
return x.CloseSend()
}
type DRPCMCPConfiguratorClient interface {
DRPCConn() drpc.Conn
GetMCPServerConfigs(ctx context.Context, in *GetMCPServerConfigsRequest) (*GetMCPServerConfigsResponse, error)
GetMCPServerAccessTokensBatch(ctx context.Context, in *GetMCPServerAccessTokensBatchRequest) (*GetMCPServerAccessTokensBatchResponse, error)
}
type drpcMCPConfiguratorClient struct {
cc drpc.Conn
}
func NewDRPCMCPConfiguratorClient(cc drpc.Conn) DRPCMCPConfiguratorClient {
return &drpcMCPConfiguratorClient{cc}
}
func (c *drpcMCPConfiguratorClient) DRPCConn() drpc.Conn { return c.cc }
func (c *drpcMCPConfiguratorClient) GetMCPServerConfigs(ctx context.Context, in *GetMCPServerConfigsRequest) (*GetMCPServerConfigsResponse, error) {
out := new(GetMCPServerConfigsResponse)
err := c.cc.Invoke(ctx, "/proto.MCPConfigurator/GetMCPServerConfigs", drpcEncoding_File_aibridged_proto_aibridged_proto{}, in, out)
if err != nil {
return nil, err
}
return out, nil
}
func (c *drpcMCPConfiguratorClient) GetMCPServerAccessTokensBatch(ctx context.Context, in *GetMCPServerAccessTokensBatchRequest) (*GetMCPServerAccessTokensBatchResponse, error) {
out := new(GetMCPServerAccessTokensBatchResponse)
err := c.cc.Invoke(ctx, "/proto.MCPConfigurator/GetMCPServerAccessTokensBatch", drpcEncoding_File_aibridged_proto_aibridged_proto{}, in, out)
if err != nil {
return nil, err
}
return out, nil
}
type DRPCMCPConfiguratorServer interface {
GetMCPServerConfigs(context.Context, *GetMCPServerConfigsRequest) (*GetMCPServerConfigsResponse, error)
GetMCPServerAccessTokensBatch(context.Context, *GetMCPServerAccessTokensBatchRequest) (*GetMCPServerAccessTokensBatchResponse, error)
}
type DRPCMCPConfiguratorUnimplementedServer struct{}
func (s *DRPCMCPConfiguratorUnimplementedServer) GetMCPServerConfigs(context.Context, *GetMCPServerConfigsRequest) (*GetMCPServerConfigsResponse, error) {
return nil, drpcerr.WithCode(errors.New("Unimplemented"), drpcerr.Unimplemented)
}
func (s *DRPCMCPConfiguratorUnimplementedServer) GetMCPServerAccessTokensBatch(context.Context, *GetMCPServerAccessTokensBatchRequest) (*GetMCPServerAccessTokensBatchResponse, error) {
return nil, drpcerr.WithCode(errors.New("Unimplemented"), drpcerr.Unimplemented)
}
type DRPCMCPConfiguratorDescription struct{}
func (DRPCMCPConfiguratorDescription) NumMethods() int { return 2 }
func (DRPCMCPConfiguratorDescription) Method(n int) (string, drpc.Encoding, drpc.Receiver, interface{}, bool) {
switch n {
case 0:
return "/proto.MCPConfigurator/GetMCPServerConfigs", drpcEncoding_File_aibridged_proto_aibridged_proto{},
func(srv interface{}, ctx context.Context, in1, in2 interface{}) (drpc.Message, error) {
return srv.(DRPCMCPConfiguratorServer).
GetMCPServerConfigs(
ctx,
in1.(*GetMCPServerConfigsRequest),
)
}, DRPCMCPConfiguratorServer.GetMCPServerConfigs, true
case 1:
return "/proto.MCPConfigurator/GetMCPServerAccessTokensBatch", drpcEncoding_File_aibridged_proto_aibridged_proto{},
func(srv interface{}, ctx context.Context, in1, in2 interface{}) (drpc.Message, error) {
return srv.(DRPCMCPConfiguratorServer).
GetMCPServerAccessTokensBatch(
ctx,
in1.(*GetMCPServerAccessTokensBatchRequest),
)
}, DRPCMCPConfiguratorServer.GetMCPServerAccessTokensBatch, true
default:
return "", nil, nil, nil, false
}
}
func DRPCRegisterMCPConfigurator(mux drpc.Mux, impl DRPCMCPConfiguratorServer) error {
return mux.Register(impl, DRPCMCPConfiguratorDescription{})
}
type DRPCMCPConfigurator_GetMCPServerConfigsStream interface {
drpc.Stream
SendAndClose(*GetMCPServerConfigsResponse) error
}
type drpcMCPConfigurator_GetMCPServerConfigsStream struct {
drpc.Stream
}
func (x *drpcMCPConfigurator_GetMCPServerConfigsStream) SendAndClose(m *GetMCPServerConfigsResponse) error {
if err := x.MsgSend(m, drpcEncoding_File_aibridged_proto_aibridged_proto{}); err != nil {
return err
}
return x.CloseSend()
}
type DRPCMCPConfigurator_GetMCPServerAccessTokensBatchStream interface {
drpc.Stream
SendAndClose(*GetMCPServerAccessTokensBatchResponse) error
}
type drpcMCPConfigurator_GetMCPServerAccessTokensBatchStream struct {
drpc.Stream
}
func (x *drpcMCPConfigurator_GetMCPServerAccessTokensBatchStream) SendAndClose(m *GetMCPServerAccessTokensBatchResponse) error {
if err := x.MsgSend(m, drpcEncoding_File_aibridged_proto_aibridged_proto{}); err != nil {
return err
}
return x.CloseSend()
}
type DRPCAuthorizerClient interface {
DRPCConn() drpc.Conn
IsAuthorized(ctx context.Context, in *IsAuthorizedRequest) (*IsAuthorizedResponse, error)
}
type drpcAuthorizerClient struct {
cc drpc.Conn
}
func NewDRPCAuthorizerClient(cc drpc.Conn) DRPCAuthorizerClient {
return &drpcAuthorizerClient{cc}
}
func (c *drpcAuthorizerClient) DRPCConn() drpc.Conn { return c.cc }
func (c *drpcAuthorizerClient) IsAuthorized(ctx context.Context, in *IsAuthorizedRequest) (*IsAuthorizedResponse, error) {
out := new(IsAuthorizedResponse)
err := c.cc.Invoke(ctx, "/proto.Authorizer/IsAuthorized", drpcEncoding_File_aibridged_proto_aibridged_proto{}, in, out)
if err != nil {
return nil, err
}
return out, nil
}
type DRPCAuthorizerServer interface {
IsAuthorized(context.Context, *IsAuthorizedRequest) (*IsAuthorizedResponse, error)
}
type DRPCAuthorizerUnimplementedServer struct{}
func (s *DRPCAuthorizerUnimplementedServer) IsAuthorized(context.Context, *IsAuthorizedRequest) (*IsAuthorizedResponse, error) {
return nil, drpcerr.WithCode(errors.New("Unimplemented"), drpcerr.Unimplemented)
}
type DRPCAuthorizerDescription struct{}
func (DRPCAuthorizerDescription) NumMethods() int { return 1 }
func (DRPCAuthorizerDescription) Method(n int) (string, drpc.Encoding, drpc.Receiver, interface{}, bool) {
switch n {
case 0:
return "/proto.Authorizer/IsAuthorized", drpcEncoding_File_aibridged_proto_aibridged_proto{},
func(srv interface{}, ctx context.Context, in1, in2 interface{}) (drpc.Message, error) {
return srv.(DRPCAuthorizerServer).
IsAuthorized(
ctx,
in1.(*IsAuthorizedRequest),
)
}, DRPCAuthorizerServer.IsAuthorized, true
default:
return "", nil, nil, nil, false
}
}
func DRPCRegisterAuthorizer(mux drpc.Mux, impl DRPCAuthorizerServer) error {
return mux.Register(impl, DRPCAuthorizerDescription{})
}
type DRPCAuthorizer_IsAuthorizedStream interface {
drpc.Stream
SendAndClose(*IsAuthorizedResponse) error
}
type drpcAuthorizer_IsAuthorizedStream struct {
drpc.Stream
}
func (x *drpcAuthorizer_IsAuthorizedStream) SendAndClose(m *IsAuthorizedResponse) error {
if err := x.MsgSend(m, drpcEncoding_File_aibridged_proto_aibridged_proto{}); err != nil {
return err
}
return x.CloseSend()
}