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:
Danny Kopping
2026-01-14 17:26:08 +02:00
committed by GitHub
parent 8d6a202ee4
commit 7d5cd06f83
13 changed files with 438 additions and 111 deletions
+4
View File
@@ -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
View File
@@ -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.
+3
View File
@@ -11970,6 +11970,9 @@ const docTemplate = `{
},
"retention": {
"type": "integer"
},
"structured_logging": {
"type": "boolean"
}
}
},
+3
View File
@@ -10631,6 +10631,9 @@
},
"retention": {
"type": "integer"
},
"structured_logging": {
"type": "boolean"
}
}
},
+11
View File
@@ -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 {
+2 -1
View File
@@ -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,
+9 -4
View File
@@ -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,
+11
View File
@@ -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
| | |
+106 -18
View File
@@ -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)
}
}
})
}
}
+4
View File
@@ -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.
+1
View File
@@ -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