mirror of
https://github.com/coder/coder.git
synced 2026-06-02 20:48:20 +00:00
feat: add aibridge structured logging (#21492)
Closes https://github.com/coder/internal/issues/1151 Sample: ``` [API] 2026-01-13 15:50:20.795 [info] coderd.aibridgedserver: interception started trace=8bb5a1d8eb10526cc46ad90f191bb468 span=a3e5b5da9546032a record_type=interception_start interception_id=97461880-4a6c-47c1-8292-3588dd715312 initiator_id=360c6167-a93a-4442-9c3e-f87a6d1cfb66 api_key_id=vg1sbUv97d provider=anthropic model=claude-opus-4-5-20251101 started_at="2026-01-13T15:50:20.790690781Z" metadata={} [API] 2026-01-13 15:50:23.741 [info] coderd.aibridgedserver: token usage recorded trace=8bb5a1d8eb10526cc46ad90f191bb468 span=a114f0cc3047296e record_type=token_usage interception_id=97461880-4a6c-47c1-8292-3588dd715312 msg_id=msg_01VJH1rYKspfun8BW29CrYEu input_tokens=10 output_tokens=8 created_at="2026-01-13T15:50:23.731587038Z" metadata={"cache_creation_input":53194,"cache_ephemeral_1h_input":0,"cache_ephemeral_5m_input":53194,"cache_read_input":0,"web_search_requests":0} [API] 2026-01-13 15:50:26.265 [info] coderd.aibridgedserver: token usage recorded trace=8bb5a1d8eb10526cc46ad90f191bb468 span=dbdafb563bff2c9c record_type=token_usage interception_id=97461880-4a6c-47c1-8292-3588dd715312 msg_id=msg_01VJH1rYKspfun8BW29CrYEu input_tokens=0 output_tokens=130 created_at="2026-01-13T15:50:26.254467904Z" metadata={} [API] 2026-01-13 15:50:26.268 [info] coderd.aibridgedserver: prompt usage recorded trace=8bb5a1d8eb10526cc46ad90f191bb468 span=da51887a757226fc record_type=prompt_usage interception_id=97461880-4a6c-47c1-8292-3588dd715312 msg_id=msg_01VJH1rYKspfun8BW29CrYEu prompt="list the jmia share price" created_at="2026-01-13T15:50:26.255299811Z" metadata={} [API] 2026-01-13 15:50:26.268 [info] coderd.aibridgedserver: interception ended trace=8bb5a1d8eb10526cc46ad90f191bb468 span=3fa25397705ee7c9 record_type=interception_end interception_id=97461880-4a6c-47c1-8292-3588dd715312 ended_at="2026-01-13T15:50:26.25555547Z" [API] 2026-01-13 15:50:26.269 [info] coderd.aibridgedserver: tool usage recorded trace=8bb5a1d8eb10526cc46ad90f191bb468 span=b54af90afc604d29 record_type=tool_usage interception_id=97461880-4a6c-47c1-8292-3588dd715312 msg_id=msg_01VJH1rYKspfun8BW29CrYEu tool=mcp__stonks__getStockPriceSnapshot input="{\"ticker\":\"JMIA\"}" server_url="" injected=false invocation_error="" created_at="2026-01-13T15:50:26.255164652Z" metadata={} ``` Structured logging is only enabled when `CODER_AIBRIDGE_STRUCTURED_LOGGING=true`. --------- Signed-off-by: Danny Kopping <danny@coder.com>
This commit is contained in:
+4
@@ -147,6 +147,10 @@ AI BRIDGE OPTIONS:
|
||||
Maximum number of AI Bridge requests per second per replica. Set to 0
|
||||
to disable (unlimited).
|
||||
|
||||
--aibridge-structured-logging bool, $CODER_AIBRIDGE_STRUCTURED_LOGGING (default: false)
|
||||
Emit structured logs for AI Bridge interception records. Use this for
|
||||
exporting these records to external SIEM or observability systems.
|
||||
|
||||
AI BRIDGE PROXY OPTIONS:
|
||||
--aibridge-proxy-cert-file string, $CODER_AIBRIDGE_PROXY_CERT_FILE
|
||||
Path to the CA certificate file for AI Bridge Proxy.
|
||||
|
||||
+4
@@ -773,6 +773,10 @@ aibridge:
|
||||
# (unlimited).
|
||||
# (default: 0, type: int)
|
||||
rateLimit: 0
|
||||
# Emit structured logs for AI Bridge interception records. Use this for exporting
|
||||
# these records to external SIEM or observability systems.
|
||||
# (default: false, type: bool)
|
||||
structuredLogging: false
|
||||
aibridgeproxy:
|
||||
# Enable the AI Bridge MITM Proxy for intercepting and decrypting AI provider
|
||||
# requests.
|
||||
|
||||
Generated
+3
@@ -11970,6 +11970,9 @@ const docTemplate = `{
|
||||
},
|
||||
"retention": {
|
||||
"type": "integer"
|
||||
},
|
||||
"structured_logging": {
|
||||
"type": "boolean"
|
||||
}
|
||||
}
|
||||
},
|
||||
|
||||
Generated
+3
@@ -10631,6 +10631,9 @@
|
||||
},
|
||||
"retention": {
|
||||
"type": "integer"
|
||||
},
|
||||
"structured_logging": {
|
||||
"type": "boolean"
|
||||
}
|
||||
}
|
||||
},
|
||||
|
||||
@@ -3484,6 +3484,16 @@ Write out the current server config as YAML to stdout.`,
|
||||
Group: &deploymentGroupAIBridge,
|
||||
YAML: "rateLimit",
|
||||
},
|
||||
{
|
||||
Name: "AI Bridge Structured Logging",
|
||||
Description: "Emit structured logs for AI Bridge interception records. Use this for exporting these records to external SIEM or observability systems.",
|
||||
Flag: "aibridge-structured-logging",
|
||||
Env: "CODER_AIBRIDGE_STRUCTURED_LOGGING",
|
||||
Value: &c.AI.BridgeConfig.StructuredLogging,
|
||||
Default: "false",
|
||||
Group: &deploymentGroupAIBridge,
|
||||
YAML: "structuredLogging",
|
||||
},
|
||||
|
||||
// AI Bridge Proxy Options
|
||||
{
|
||||
@@ -3610,6 +3620,7 @@ type AIBridgeConfig struct {
|
||||
Retention serpent.Duration `json:"retention" typescript:",notnull"`
|
||||
MaxConcurrency serpent.Int64 `json:"max_concurrency" typescript:",notnull"`
|
||||
RateLimit serpent.Int64 `json:"rate_limit" typescript:",notnull"`
|
||||
StructuredLogging serpent.Bool `json:"structured_logging" typescript:",notnull"`
|
||||
}
|
||||
|
||||
type AIBridgeOpenAIConfig struct {
|
||||
|
||||
Generated
+2
-1
@@ -191,7 +191,8 @@ curl -X GET http://coder-server:8080/api/v2/deployment/config \
|
||||
"key": "string"
|
||||
},
|
||||
"rate_limit": 0,
|
||||
"retention": 0
|
||||
"retention": 0,
|
||||
"structured_logging": true
|
||||
}
|
||||
},
|
||||
"allow_workspace_renames": true,
|
||||
|
||||
Generated
+9
-4
@@ -396,7 +396,8 @@
|
||||
"key": "string"
|
||||
},
|
||||
"rate_limit": 0,
|
||||
"retention": 0
|
||||
"retention": 0,
|
||||
"structured_logging": true
|
||||
}
|
||||
```
|
||||
|
||||
@@ -412,6 +413,7 @@
|
||||
| `openai` | [codersdk.AIBridgeOpenAIConfig](#codersdkaibridgeopenaiconfig) | false | | |
|
||||
| `rate_limit` | integer | false | | |
|
||||
| `retention` | integer | false | | |
|
||||
| `structured_logging` | boolean | false | | |
|
||||
|
||||
## codersdk.AIBridgeInterception
|
||||
|
||||
@@ -743,7 +745,8 @@
|
||||
"key": "string"
|
||||
},
|
||||
"rate_limit": 0,
|
||||
"retention": 0
|
||||
"retention": 0,
|
||||
"structured_logging": true
|
||||
}
|
||||
}
|
||||
```
|
||||
@@ -2658,7 +2661,8 @@ CreateWorkspaceRequest provides options for creating a new workspace. Only one o
|
||||
"key": "string"
|
||||
},
|
||||
"rate_limit": 0,
|
||||
"retention": 0
|
||||
"retention": 0,
|
||||
"structured_logging": true
|
||||
}
|
||||
},
|
||||
"allow_workspace_renames": true,
|
||||
@@ -3202,7 +3206,8 @@ CreateWorkspaceRequest provides options for creating a new workspace. Only one o
|
||||
"key": "string"
|
||||
},
|
||||
"rate_limit": 0,
|
||||
"retention": 0
|
||||
"retention": 0,
|
||||
"structured_logging": true
|
||||
}
|
||||
},
|
||||
"allow_workspace_renames": true,
|
||||
|
||||
Generated
+11
@@ -1836,6 +1836,17 @@ Maximum number of concurrent AI Bridge requests per replica. Set to 0 to disable
|
||||
|
||||
Maximum number of AI Bridge requests per second per replica. Set to 0 to disable (unlimited).
|
||||
|
||||
### --aibridge-structured-logging
|
||||
|
||||
| | |
|
||||
|-------------|-------------------------------------------------|
|
||||
| Type | <code>bool</code> |
|
||||
| Environment | <code>$CODER_AIBRIDGE_STRUCTURED_LOGGING</code> |
|
||||
| YAML | <code>aibridge.structuredLogging</code> |
|
||||
| Default | <code>false</code> |
|
||||
|
||||
Emit structured logs for AI Bridge interception records. Use this for exporting these records to external SIEM or observability systems.
|
||||
|
||||
### --aibridge-proxy-enabled
|
||||
|
||||
| | |
|
||||
|
||||
@@ -46,6 +46,10 @@ var (
|
||||
ErrNoExternalAuthLinkFound = xerrors.New("no external auth link found")
|
||||
)
|
||||
|
||||
const (
|
||||
InterceptionLogMarker = "interception log"
|
||||
)
|
||||
|
||||
var _ aibridged.DRPCServer = &Server{}
|
||||
|
||||
type store interface {
|
||||
@@ -73,7 +77,8 @@ type Server struct {
|
||||
logger slog.Logger
|
||||
externalAuthConfigs map[string]*externalauth.Config
|
||||
|
||||
coderMCPConfig *proto.MCPServerConfig // may be nil if not available
|
||||
coderMCPConfig *proto.MCPServerConfig // may be nil if not available
|
||||
structuredLogging bool
|
||||
}
|
||||
|
||||
func NewServer(lifecycleCtx context.Context, store store, logger slog.Logger, accessURL string,
|
||||
@@ -92,8 +97,9 @@ func NewServer(lifecycleCtx context.Context, store store, logger slog.Logger, ac
|
||||
srv := &Server{
|
||||
lifecycleCtx: lifecycleCtx,
|
||||
store: store,
|
||||
logger: logger.Named("aibridgedserver"),
|
||||
logger: logger,
|
||||
externalAuthConfigs: eac,
|
||||
structuredLogging: bridgeCfg.StructuredLogging.Value(),
|
||||
}
|
||||
|
||||
if bridgeCfg.InjectCoderMCPTools {
|
||||
@@ -123,13 +129,33 @@ func (s *Server) RecordInterception(ctx context.Context, in *proto.RecordInterce
|
||||
return nil, xerrors.Errorf("empty API key ID")
|
||||
}
|
||||
|
||||
metadata := metadataToMap(in.GetMetadata())
|
||||
|
||||
if s.structuredLogging {
|
||||
s.logger.Info(ctx, InterceptionLogMarker,
|
||||
slog.F("record_type", "interception_start"),
|
||||
slog.F("interception_id", intcID.String()),
|
||||
slog.F("initiator_id", initID.String()),
|
||||
slog.F("api_key_id", in.ApiKeyId),
|
||||
slog.F("provider", in.Provider),
|
||||
slog.F("model", in.Model),
|
||||
slog.F("started_at", in.StartedAt.AsTime()),
|
||||
slog.F("metadata", metadata),
|
||||
)
|
||||
}
|
||||
|
||||
out, err := json.Marshal(metadata)
|
||||
if err != nil {
|
||||
s.logger.Warn(ctx, "failed to marshal aibridge metadata from proto to JSON", slog.F("metadata", in), slog.Error(err))
|
||||
}
|
||||
|
||||
_, err = s.store.InsertAIBridgeInterception(ctx, database.InsertAIBridgeInterceptionParams{
|
||||
ID: intcID,
|
||||
APIKeyID: sql.NullString{String: in.ApiKeyId, Valid: true},
|
||||
InitiatorID: initID,
|
||||
Provider: in.Provider,
|
||||
Model: in.Model,
|
||||
Metadata: marshalMetadata(ctx, s.logger, in.GetMetadata()),
|
||||
Metadata: out,
|
||||
StartedAt: in.StartedAt.AsTime(),
|
||||
})
|
||||
if err != nil {
|
||||
@@ -148,6 +174,14 @@ func (s *Server) RecordInterceptionEnded(ctx context.Context, in *proto.RecordIn
|
||||
return nil, xerrors.Errorf("invalid interception ID %q: %w", in.GetId(), err)
|
||||
}
|
||||
|
||||
if s.structuredLogging {
|
||||
s.logger.Info(ctx, InterceptionLogMarker,
|
||||
slog.F("record_type", "interception_end"),
|
||||
slog.F("interception_id", intcID.String()),
|
||||
slog.F("ended_at", in.EndedAt.AsTime()),
|
||||
)
|
||||
}
|
||||
|
||||
_, err = s.store.UpdateAIBridgeInterceptionEnded(ctx, database.UpdateAIBridgeInterceptionEndedParams{
|
||||
ID: intcID,
|
||||
EndedAt: in.EndedAt.AsTime(),
|
||||
@@ -168,18 +202,38 @@ func (s *Server) RecordTokenUsage(ctx context.Context, in *proto.RecordTokenUsag
|
||||
return nil, xerrors.Errorf("failed to parse interception_id %q: %w", in.GetInterceptionId(), err)
|
||||
}
|
||||
|
||||
metadata := metadataToMap(in.GetMetadata())
|
||||
|
||||
if s.structuredLogging {
|
||||
s.logger.Info(ctx, InterceptionLogMarker,
|
||||
slog.F("record_type", "token_usage"),
|
||||
slog.F("interception_id", intcID.String()),
|
||||
slog.F("msg_id", in.GetMsgId()),
|
||||
slog.F("input_tokens", in.GetInputTokens()),
|
||||
slog.F("output_tokens", in.GetOutputTokens()),
|
||||
slog.F("created_at", in.GetCreatedAt().AsTime()),
|
||||
slog.F("metadata", metadata),
|
||||
)
|
||||
}
|
||||
|
||||
out, err := json.Marshal(metadata)
|
||||
if err != nil {
|
||||
s.logger.Warn(ctx, "failed to marshal aibridge metadata from proto to JSON", slog.F("metadata", in), slog.Error(err))
|
||||
}
|
||||
|
||||
_, err = s.store.InsertAIBridgeTokenUsage(ctx, database.InsertAIBridgeTokenUsageParams{
|
||||
ID: uuid.New(),
|
||||
InterceptionID: intcID,
|
||||
ProviderResponseID: in.GetMsgId(),
|
||||
InputTokens: in.GetInputTokens(),
|
||||
OutputTokens: in.GetOutputTokens(),
|
||||
Metadata: marshalMetadata(ctx, s.logger, in.GetMetadata()),
|
||||
Metadata: out,
|
||||
CreatedAt: in.GetCreatedAt().AsTime(),
|
||||
})
|
||||
if err != nil {
|
||||
return nil, xerrors.Errorf("insert token usage: %w", err)
|
||||
}
|
||||
|
||||
return &proto.RecordTokenUsageResponse{}, nil
|
||||
}
|
||||
|
||||
@@ -192,17 +246,36 @@ func (s *Server) RecordPromptUsage(ctx context.Context, in *proto.RecordPromptUs
|
||||
return nil, xerrors.Errorf("failed to parse interception_id %q: %w", in.GetInterceptionId(), err)
|
||||
}
|
||||
|
||||
metadata := metadataToMap(in.GetMetadata())
|
||||
|
||||
if s.structuredLogging {
|
||||
s.logger.Info(ctx, InterceptionLogMarker,
|
||||
slog.F("record_type", "prompt_usage"),
|
||||
slog.F("interception_id", intcID.String()),
|
||||
slog.F("msg_id", in.GetMsgId()),
|
||||
slog.F("prompt", in.GetPrompt()),
|
||||
slog.F("created_at", in.GetCreatedAt().AsTime()),
|
||||
slog.F("metadata", metadata),
|
||||
)
|
||||
}
|
||||
|
||||
out, err := json.Marshal(metadata)
|
||||
if err != nil {
|
||||
s.logger.Warn(ctx, "failed to marshal aibridge metadata from proto to JSON", slog.F("metadata", in), slog.Error(err))
|
||||
}
|
||||
|
||||
_, err = s.store.InsertAIBridgeUserPrompt(ctx, database.InsertAIBridgeUserPromptParams{
|
||||
ID: uuid.New(),
|
||||
InterceptionID: intcID,
|
||||
ProviderResponseID: in.GetMsgId(),
|
||||
Prompt: in.GetPrompt(),
|
||||
Metadata: marshalMetadata(ctx, s.logger, in.GetMetadata()),
|
||||
Metadata: out,
|
||||
CreatedAt: in.GetCreatedAt().AsTime(),
|
||||
})
|
||||
if err != nil {
|
||||
return nil, xerrors.Errorf("insert user prompt: %w", err)
|
||||
}
|
||||
|
||||
return &proto.RecordPromptUsageResponse{}, nil
|
||||
}
|
||||
|
||||
@@ -215,6 +288,28 @@ func (s *Server) RecordToolUsage(ctx context.Context, in *proto.RecordToolUsageR
|
||||
return nil, xerrors.Errorf("failed to parse interception_id %q: %w", in.GetInterceptionId(), err)
|
||||
}
|
||||
|
||||
metadata := metadataToMap(in.GetMetadata())
|
||||
|
||||
if s.structuredLogging {
|
||||
s.logger.Info(ctx, InterceptionLogMarker,
|
||||
slog.F("record_type", "tool_usage"),
|
||||
slog.F("interception_id", intcID.String()),
|
||||
slog.F("msg_id", in.GetMsgId()),
|
||||
slog.F("tool", in.GetTool()),
|
||||
slog.F("input", in.GetInput()),
|
||||
slog.F("server_url", in.GetServerUrl()),
|
||||
slog.F("injected", in.GetInjected()),
|
||||
slog.F("invocation_error", in.GetInvocationError()),
|
||||
slog.F("created_at", in.GetCreatedAt().AsTime()),
|
||||
slog.F("metadata", metadata),
|
||||
)
|
||||
}
|
||||
|
||||
out, err := json.Marshal(metadata)
|
||||
if err != nil {
|
||||
s.logger.Warn(ctx, "failed to marshal aibridge metadata from proto to JSON", slog.F("metadata", in), slog.Error(err))
|
||||
}
|
||||
|
||||
_, err = s.store.InsertAIBridgeToolUsage(ctx, database.InsertAIBridgeToolUsageParams{
|
||||
ID: uuid.New(),
|
||||
InterceptionID: intcID,
|
||||
@@ -224,12 +319,13 @@ func (s *Server) RecordToolUsage(ctx context.Context, in *proto.RecordToolUsageR
|
||||
Input: in.GetInput(),
|
||||
Injected: in.GetInjected(),
|
||||
InvocationError: sql.NullString{String: in.GetInvocationError(), Valid: in.InvocationError != nil},
|
||||
Metadata: marshalMetadata(ctx, s.logger, in.GetMetadata()),
|
||||
Metadata: out,
|
||||
CreatedAt: in.GetCreatedAt().AsTime(),
|
||||
})
|
||||
if err != nil {
|
||||
return nil, xerrors.Errorf("insert tool usage: %w", err)
|
||||
}
|
||||
|
||||
return &proto.RecordToolUsageResponse{}, nil
|
||||
}
|
||||
|
||||
@@ -433,24 +529,16 @@ func getCoderMCPServerConfig(experiments codersdk.Experiments, accessURL string)
|
||||
}, nil
|
||||
}
|
||||
|
||||
// marshalMetadata attempts to marshal the given metadata map into a
|
||||
// JSON-encoded byte slice. If the marshaling fails, the function logs a
|
||||
// warning and returns nil. The supplied context is only used for logging.
|
||||
func marshalMetadata(ctx context.Context, logger slog.Logger, in map[string]*anypb.Any) []byte {
|
||||
mdMap := make(map[string]any, len(in))
|
||||
func metadataToMap(in map[string]*anypb.Any) map[string]any {
|
||||
meta := make(map[string]any, len(in))
|
||||
for k, v := range in {
|
||||
if v == nil {
|
||||
continue
|
||||
}
|
||||
var sv structpb.Value
|
||||
if err := v.UnmarshalTo(&sv); err == nil {
|
||||
mdMap[k] = sv.AsInterface()
|
||||
meta[k] = sv.AsInterface()
|
||||
}
|
||||
}
|
||||
out, err := json.Marshal(mdMap)
|
||||
if err != nil {
|
||||
logger.Warn(ctx, "failed to marshal aibridge metadata from proto to JSON", slog.F("metadata", in), slog.Error(err))
|
||||
return nil
|
||||
}
|
||||
return out
|
||||
return meta
|
||||
}
|
||||
|
||||
@@ -1,88 +0,0 @@
|
||||
package aibridgedserver
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"math"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
"google.golang.org/protobuf/proto"
|
||||
"google.golang.org/protobuf/types/known/anypb"
|
||||
"google.golang.org/protobuf/types/known/structpb"
|
||||
|
||||
"cdr.dev/slog/v3"
|
||||
"cdr.dev/slog/v3/sloggers/slogtest"
|
||||
)
|
||||
|
||||
func TestMarshalMetadata(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
t.Run("NilData", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).Leveled(slog.LevelDebug)
|
||||
out := marshalMetadata(context.Background(), logger, nil)
|
||||
require.JSONEq(t, "{}", string(out))
|
||||
})
|
||||
|
||||
t.Run("WithData", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).Leveled(slog.LevelDebug)
|
||||
|
||||
list := structpb.NewListValue(&structpb.ListValue{Values: []*structpb.Value{
|
||||
structpb.NewStringValue("a"),
|
||||
structpb.NewNumberValue(1),
|
||||
structpb.NewBoolValue(false),
|
||||
}})
|
||||
obj := structpb.NewStructValue(&structpb.Struct{Fields: map[string]*structpb.Value{
|
||||
"a": structpb.NewStringValue("b"),
|
||||
"n": structpb.NewNumberValue(3),
|
||||
}})
|
||||
|
||||
nonValue := mustMarshalAny(t, &structpb.Struct{Fields: map[string]*structpb.Value{
|
||||
"ignored": structpb.NewStringValue("yes"),
|
||||
}})
|
||||
invalid := &anypb.Any{TypeUrl: "type.googleapis.com/google.protobuf.Value", Value: []byte{0xff, 0x00}}
|
||||
|
||||
in := map[string]*anypb.Any{
|
||||
"null": mustMarshalAny(t, structpb.NewNullValue()),
|
||||
// Scalars
|
||||
"string": mustMarshalAny(t, structpb.NewStringValue("hello")),
|
||||
"bool": mustMarshalAny(t, structpb.NewBoolValue(true)),
|
||||
"number": mustMarshalAny(t, structpb.NewNumberValue(42)),
|
||||
// Complex types
|
||||
"list": mustMarshalAny(t, list),
|
||||
"object": mustMarshalAny(t, obj),
|
||||
// Extra valid entries
|
||||
"ok": mustMarshalAny(t, structpb.NewStringValue("present")),
|
||||
"nan": mustMarshalAny(t, structpb.NewNumberValue(math.NaN())),
|
||||
// Entries that should be ignored
|
||||
"invalid": invalid,
|
||||
"non_value": nonValue,
|
||||
}
|
||||
|
||||
out := marshalMetadata(context.Background(), logger, in)
|
||||
require.NotNil(t, out)
|
||||
var got map[string]any
|
||||
require.NoError(t, json.Unmarshal(out, &got))
|
||||
|
||||
expected := map[string]any{
|
||||
"string": "hello",
|
||||
"bool": true,
|
||||
"number": float64(42),
|
||||
"null": nil,
|
||||
"list": []any{"a", float64(1), false},
|
||||
"object": map[string]any{"a": "b", "n": float64(3)},
|
||||
"ok": "present",
|
||||
"nan": "NaN",
|
||||
}
|
||||
require.Equal(t, expected, got)
|
||||
})
|
||||
}
|
||||
|
||||
func mustMarshalAny(t testing.TB, m proto.Message) *anypb.Any {
|
||||
t.Helper()
|
||||
a, err := anypb.New(m)
|
||||
require.NoError(t, err)
|
||||
return a
|
||||
}
|
||||
@@ -1,6 +1,8 @@
|
||||
package aibridgedserver_test
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
@@ -20,6 +22,8 @@ import (
|
||||
"google.golang.org/protobuf/types/known/structpb"
|
||||
"google.golang.org/protobuf/types/known/timestamppb"
|
||||
|
||||
"cdr.dev/slog/v3"
|
||||
"cdr.dev/slog/v3/sloggers/slogjson"
|
||||
"github.com/coder/coder/v2/coderd/apikey"
|
||||
"github.com/coder/coder/v2/coderd/database"
|
||||
"github.com/coder/coder/v2/coderd/database/dbmock"
|
||||
@@ -832,3 +836,279 @@ func mustMarshalAny(t *testing.T, msg protobufproto.Message) *anypb.Any {
|
||||
func strPtr(s string) *string {
|
||||
return &s
|
||||
}
|
||||
|
||||
// logLine represents a parsed JSON log entry.
|
||||
type logLine struct {
|
||||
Msg string `json:"msg"`
|
||||
Level string `json:"level"`
|
||||
Fields map[string]any `json:"fields"`
|
||||
}
|
||||
|
||||
// parseLogLines parses JSON log lines from a buffer.
|
||||
func parseLogLines(buf *bytes.Buffer) []logLine {
|
||||
var lines []logLine
|
||||
scanner := bufio.NewScanner(buf)
|
||||
for scanner.Scan() {
|
||||
var line logLine
|
||||
if err := json.Unmarshal(scanner.Bytes(), &line); err == nil {
|
||||
lines = append(lines, line)
|
||||
}
|
||||
}
|
||||
return lines
|
||||
}
|
||||
|
||||
// getLogLinesWithMessage returns all log lines with the given message.
|
||||
func getLogLinesWithMessage(lines []logLine, msg string) []logLine {
|
||||
var result []logLine
|
||||
for _, line := range lines {
|
||||
if line.Msg == msg {
|
||||
result = append(result, line)
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
func TestStructuredLogging(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
metadataProto := map[string]*anypb.Any{
|
||||
"key": mustMarshalAny(t, &structpb.Value{Kind: &structpb.Value_StringValue{StringValue: "value"}}),
|
||||
}
|
||||
|
||||
type testCase struct {
|
||||
name string
|
||||
structuredLogging bool
|
||||
expectedErr error
|
||||
setupMocks func(db *dbmock.MockStore, interceptionID uuid.UUID)
|
||||
recordFn func(srv *aibridgedserver.Server, ctx context.Context, interceptionID uuid.UUID) error
|
||||
expectedFields map[string]any
|
||||
}
|
||||
|
||||
interceptionID := uuid.UUID{1}
|
||||
initiatorID := uuid.UUID{2}
|
||||
|
||||
cases := []testCase{
|
||||
{
|
||||
name: "RecordInterception_logs_when_enabled",
|
||||
structuredLogging: true,
|
||||
setupMocks: func(db *dbmock.MockStore, intcID uuid.UUID) {
|
||||
db.EXPECT().InsertAIBridgeInterception(gomock.Any(), gomock.Any()).Return(database.AIBridgeInterception{
|
||||
ID: intcID,
|
||||
InitiatorID: initiatorID,
|
||||
}, nil)
|
||||
},
|
||||
recordFn: func(srv *aibridgedserver.Server, ctx context.Context, intcID uuid.UUID) error {
|
||||
_, err := srv.RecordInterception(ctx, &proto.RecordInterceptionRequest{
|
||||
Id: intcID.String(),
|
||||
ApiKeyId: "api-key-123",
|
||||
InitiatorId: initiatorID.String(),
|
||||
Provider: "anthropic",
|
||||
Model: "claude-4-opus",
|
||||
Metadata: metadataProto,
|
||||
StartedAt: timestamppb.Now(),
|
||||
})
|
||||
return err
|
||||
},
|
||||
expectedFields: map[string]any{
|
||||
"record_type": "interception_start",
|
||||
"interception_id": interceptionID.String(),
|
||||
"initiator_id": initiatorID.String(),
|
||||
"provider": "anthropic",
|
||||
"model": "claude-4-opus",
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "RecordInterception_does_not_log_when_disabled",
|
||||
structuredLogging: false,
|
||||
setupMocks: func(db *dbmock.MockStore, intcID uuid.UUID) {
|
||||
db.EXPECT().InsertAIBridgeInterception(gomock.Any(), gomock.Any()).Return(database.AIBridgeInterception{
|
||||
ID: intcID,
|
||||
InitiatorID: initiatorID,
|
||||
}, nil)
|
||||
},
|
||||
recordFn: func(srv *aibridgedserver.Server, ctx context.Context, intcID uuid.UUID) error {
|
||||
_, err := srv.RecordInterception(ctx, &proto.RecordInterceptionRequest{
|
||||
Id: intcID.String(),
|
||||
ApiKeyId: "api-key-123",
|
||||
InitiatorId: initiatorID.String(),
|
||||
Provider: "anthropic",
|
||||
Model: "claude-4-opus",
|
||||
StartedAt: timestamppb.Now(),
|
||||
})
|
||||
return err
|
||||
},
|
||||
expectedFields: nil, // No log expected.
|
||||
},
|
||||
{
|
||||
name: "RecordInterception_log_on_db_error",
|
||||
structuredLogging: true,
|
||||
expectedErr: sql.ErrConnDone,
|
||||
setupMocks: func(db *dbmock.MockStore, intcID uuid.UUID) {
|
||||
db.EXPECT().InsertAIBridgeInterception(gomock.Any(), gomock.Any()).Return(database.AIBridgeInterception{}, sql.ErrConnDone)
|
||||
},
|
||||
recordFn: func(srv *aibridgedserver.Server, ctx context.Context, intcID uuid.UUID) error {
|
||||
_, err := srv.RecordInterception(ctx, &proto.RecordInterceptionRequest{
|
||||
Id: intcID.String(),
|
||||
ApiKeyId: "api-key-123",
|
||||
InitiatorId: initiatorID.String(),
|
||||
Provider: "anthropic",
|
||||
Model: "claude-4-opus",
|
||||
StartedAt: timestamppb.Now(),
|
||||
})
|
||||
return err
|
||||
},
|
||||
// Even though the database call errored, we must still write the logs.
|
||||
expectedFields: map[string]any{
|
||||
"record_type": "interception_start",
|
||||
"interception_id": interceptionID.String(),
|
||||
"initiator_id": initiatorID.String(),
|
||||
"provider": "anthropic",
|
||||
"model": "claude-4-opus",
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "RecordInterceptionEnded_logs_when_enabled",
|
||||
structuredLogging: true,
|
||||
setupMocks: func(db *dbmock.MockStore, intcID uuid.UUID) {
|
||||
db.EXPECT().UpdateAIBridgeInterceptionEnded(gomock.Any(), gomock.Any()).Return(database.AIBridgeInterception{
|
||||
ID: intcID,
|
||||
}, nil)
|
||||
},
|
||||
recordFn: func(srv *aibridgedserver.Server, ctx context.Context, intcID uuid.UUID) error {
|
||||
_, err := srv.RecordInterceptionEnded(ctx, &proto.RecordInterceptionEndedRequest{
|
||||
Id: intcID.String(),
|
||||
EndedAt: timestamppb.Now(),
|
||||
})
|
||||
return err
|
||||
},
|
||||
expectedFields: map[string]any{
|
||||
"record_type": "interception_end",
|
||||
"interception_id": interceptionID.String(),
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "RecordTokenUsage_logs_when_enabled",
|
||||
structuredLogging: true,
|
||||
setupMocks: func(db *dbmock.MockStore, intcID uuid.UUID) {
|
||||
db.EXPECT().InsertAIBridgeTokenUsage(gomock.Any(), gomock.Any()).Return(database.AIBridgeTokenUsage{
|
||||
ID: uuid.New(),
|
||||
InterceptionID: intcID,
|
||||
}, nil)
|
||||
},
|
||||
recordFn: func(srv *aibridgedserver.Server, ctx context.Context, intcID uuid.UUID) error {
|
||||
_, err := srv.RecordTokenUsage(ctx, &proto.RecordTokenUsageRequest{
|
||||
InterceptionId: intcID.String(),
|
||||
MsgId: "msg_123",
|
||||
InputTokens: 100,
|
||||
OutputTokens: 200,
|
||||
Metadata: metadataProto,
|
||||
CreatedAt: timestamppb.Now(),
|
||||
})
|
||||
return err
|
||||
},
|
||||
expectedFields: map[string]any{
|
||||
"record_type": "token_usage",
|
||||
"interception_id": interceptionID.String(),
|
||||
"input_tokens": float64(100), // JSON numbers are float64.
|
||||
"output_tokens": float64(200),
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "RecordPromptUsage_logs_when_enabled",
|
||||
structuredLogging: true,
|
||||
setupMocks: func(db *dbmock.MockStore, intcID uuid.UUID) {
|
||||
db.EXPECT().InsertAIBridgeUserPrompt(gomock.Any(), gomock.Any()).Return(database.AIBridgeUserPrompt{
|
||||
ID: uuid.New(),
|
||||
InterceptionID: intcID,
|
||||
}, nil)
|
||||
},
|
||||
recordFn: func(srv *aibridgedserver.Server, ctx context.Context, intcID uuid.UUID) error {
|
||||
_, err := srv.RecordPromptUsage(ctx, &proto.RecordPromptUsageRequest{
|
||||
InterceptionId: intcID.String(),
|
||||
MsgId: "msg_123",
|
||||
Prompt: "Hello, Claude!",
|
||||
Metadata: metadataProto,
|
||||
CreatedAt: timestamppb.Now(),
|
||||
})
|
||||
return err
|
||||
},
|
||||
expectedFields: map[string]any{
|
||||
"record_type": "prompt_usage",
|
||||
"interception_id": interceptionID.String(),
|
||||
"prompt": "Hello, Claude!",
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "RecordToolUsage_logs_when_enabled",
|
||||
structuredLogging: true,
|
||||
setupMocks: func(db *dbmock.MockStore, intcID uuid.UUID) {
|
||||
db.EXPECT().InsertAIBridgeToolUsage(gomock.Any(), gomock.Any()).Return(database.AIBridgeToolUsage{
|
||||
ID: uuid.New(),
|
||||
InterceptionID: intcID,
|
||||
}, nil)
|
||||
},
|
||||
recordFn: func(srv *aibridgedserver.Server, ctx context.Context, intcID uuid.UUID) error {
|
||||
_, err := srv.RecordToolUsage(ctx, &proto.RecordToolUsageRequest{
|
||||
InterceptionId: intcID.String(),
|
||||
MsgId: "msg_123",
|
||||
ServerUrl: strPtr("https://api.example.com"),
|
||||
Tool: "read_file",
|
||||
Input: `{"path": "/etc/hosts"}`,
|
||||
Injected: true,
|
||||
InvocationError: strPtr("permission denied"),
|
||||
Metadata: metadataProto,
|
||||
CreatedAt: timestamppb.Now(),
|
||||
})
|
||||
return err
|
||||
},
|
||||
expectedFields: map[string]any{
|
||||
"record_type": "tool_usage",
|
||||
"interception_id": interceptionID.String(),
|
||||
"tool": "read_file",
|
||||
"input": `{"path": "/etc/hosts"}`,
|
||||
"injected": true,
|
||||
"invocation_error": "permission denied",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ctrl := gomock.NewController(t)
|
||||
db := dbmock.NewMockStore(ctrl)
|
||||
buf := &bytes.Buffer{}
|
||||
logger := slog.Make(slogjson.Sink(buf)).Leveled(slog.LevelDebug)
|
||||
|
||||
tc.setupMocks(db, interceptionID)
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
srv, err := aibridgedserver.NewServer(ctx, db, logger, "/", codersdk.AIBridgeConfig{
|
||||
StructuredLogging: serpent.Bool(tc.structuredLogging),
|
||||
}, nil, requiredExperiments)
|
||||
require.NoError(t, err)
|
||||
|
||||
err = tc.recordFn(srv, ctx, interceptionID)
|
||||
if tc.expectedErr != nil {
|
||||
require.Error(t, err)
|
||||
} else {
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
lines := parseLogLines(buf)
|
||||
if tc.expectedFields == nil {
|
||||
// No log expected (disabled or error case).
|
||||
require.Empty(t, lines)
|
||||
} else {
|
||||
matchedLines := getLogLinesWithMessage(lines, aibridgedserver.InterceptionLogMarker)
|
||||
require.Len(t, matchedLines, 1, "expected exactly one log line with message %q", aibridgedserver.InterceptionLogMarker)
|
||||
|
||||
fields := matchedLines[0].Fields
|
||||
for key, expected := range tc.expectedFields {
|
||||
require.Equal(t, expected, fields[key], "field %q mismatch", key)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -148,6 +148,10 @@ AI BRIDGE OPTIONS:
|
||||
Maximum number of AI Bridge requests per second per replica. Set to 0
|
||||
to disable (unlimited).
|
||||
|
||||
--aibridge-structured-logging bool, $CODER_AIBRIDGE_STRUCTURED_LOGGING (default: false)
|
||||
Emit structured logs for AI Bridge interception records. Use this for
|
||||
exporting these records to external SIEM or observability systems.
|
||||
|
||||
AI BRIDGE PROXY OPTIONS:
|
||||
--aibridge-proxy-cert-file string, $CODER_AIBRIDGE_PROXY_CERT_FILE
|
||||
Path to the CA certificate file for AI Bridge Proxy.
|
||||
|
||||
Generated
+1
@@ -35,6 +35,7 @@ export interface AIBridgeConfig {
|
||||
readonly retention: number;
|
||||
readonly max_concurrency: number;
|
||||
readonly rate_limit: number;
|
||||
readonly structured_logging: boolean;
|
||||
}
|
||||
|
||||
// From codersdk/aibridge.go
|
||||
|
||||
Reference in New Issue
Block a user