Thanks to visit codestin.com
Credit goes to github.com

Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 12 additions & 2 deletions aibridge/config/config.go
Original file line number Diff line number Diff line change
Expand Up @@ -70,10 +70,20 @@ type AWSBedrock struct {
Protocol BedrockProtocol
}

// ResolvedProtocol returns the configured protocol, mapping the empty value to

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit [CRF-6] The ResolvedProtocol doc says "legacy" twice in one sentence. (Gon)

"mapping the empty value to the legacy InvokeModel protocol so existing providers keep the legacy behavior." Drop the first: "returns the configured protocol, or BedrockProtocolInvokeModel when unset, so existing providers keep legacy behavior."

🤖

// the legacy InvokeModel protocol so existing providers keep the legacy
// behavior.
func (c AWSBedrock) ResolvedProtocol() BedrockProtocol {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P3 [CRF-5] ResolvedProtocol() centralizes the empty->InvokeModel mapping but is wired to only the trace attribute, leaving the same invariant hand-rolled in two other places. (Mafuuu P3, Ryosuke P3, Chopper Nit, Razor Nit)

The empty-value default now lives in three spots: here, isBedrockInvokeModel() (base.go:143, Protocol == "" || Protocol == BedrockProtocolInvokeModel), and Validate() (config.go:86, case "", BedrockProtocolInvokeModel). Behavior is identical today, so nothing breaks now. The consequence is latent divergence: if the zero-value default ever changes, a maintainer updating ResolvedProtocol() leaves the routing predicates disagreeing, so Model()/dispatch route down one protocol while the trace attribute reports the other. Route isBedrockInvokeModel()/isBedrockMantle() through ResolvedProtocol() so the invariant lives in one place.

🤖

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit [CRF-7] ResolvedProtocol collapses to return cmp.Or(c.Protocol, BedrockProtocolInvokeModel). (Ging-Go)

BedrockProtocol is a string type, so cmp.Or (Go 1.22+) returns c.Protocol when non-empty else the default, identical behavior with no side effects. go.mod is on 1.26, and cmp.Or is already used in the tree (e.g. aibridge/internal/testutil/mockupstream.go:214).

🤖

if c.Protocol == "" {
return BedrockProtocolInvokeModel
}
return c.Protocol
}

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Potentially a larger refactor, so take it or leave it: unexport c.Protocol and rename this method Protocol() so that callers are forced to use this method that handles the zero value.

Alternatively, make BedrockProtocolInvokeModel the empty string?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Potentially a larger refactor, so take it or leave it: unexport c.Protocol and rename this method Protocol() so that callers are forced to use this method that handles the zero value.

I think this requires too many changes. We'd need to add a constructor and update dozens of call sites.
It also goes against the pattern we use in config/config.go, where all fields are exported.

Alternatively, make BedrockProtocolInvokeModel the empty string?

That's an interesting idea—I hadn't considered it when I first implemented this.
At this point, though, it would require additional work, since we'd also need to update the codersdk layer.
I'm also not fully convinced it's an improvement. An empty string can be confusing at the API level.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I wouldn't consider this confusing:

type BedrockProtocol string
const (
  BedrockProtocolDefault     = ""
  BedrockProtocolInvokeModel = "invoke-model"
  BedrockProtocolMantle      = "mantle"
)

...

  switch (proto) {
    case BedrockProtocolDefault, BedrockProtocolInvokeModel:
      // handle invoke-model
    case BedrockProtocolMantle:
      // handle mantle
    default:
      // unknown, error
  }

But it's fair enough to leave it as a follow-up. We do still have a window to change the bedrock protocol stuff around before the next release though!

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I thought you mean:

type BedrockProtocol string
const (
  BedrockProtocolInvokeModel = ""
  BedrockProtocolMantle      = "mantle"
)

...

  switch (proto) {
    case BedrockProtocolInvokeModel:
      // handle invoke-model
    case BedrockProtocolMantle:
      // handle mantle
    default:
      // unknown, error
  }


// Validate verifies protocol-specific Bedrock configuration.
func (c AWSBedrock) Validate() error {
switch c.Protocol {
case "", BedrockProtocolInvokeModel:
switch c.ResolvedProtocol() {
case BedrockProtocolInvokeModel:
if c.Region == "" && c.BaseURL == "" {
return xerrors.New("region or base url required")
}
Expand Down
11 changes: 7 additions & 4 deletions aibridge/intercept/messages/base.go
Original file line number Diff line number Diff line change
Expand Up @@ -133,14 +133,13 @@ func (i *interceptionBase) CorrelatingToolCallID() *string {
// isBedrockMantle reports whether the interception targets the Bedrock mantle
// protocol.
func (i *interceptionBase) isBedrockMantle() bool {
return i.bedrock != nil && i.bedrock.Cfg.Protocol == aibconfig.BedrockProtocolMantle
return i.bedrock != nil && i.bedrock.Cfg.ResolvedProtocol() == aibconfig.BedrockProtocolMantle
}

// isBedrockInvokeModel reports whether the interception targets the Bedrock
// InvokeModel protocol.
func (i *interceptionBase) isBedrockInvokeModel() bool {
return i.bedrock != nil &&
(i.bedrock.Cfg.Protocol == "" || i.bedrock.Cfg.Protocol == aibconfig.BedrockProtocolInvokeModel)
return i.bedrock != nil && i.bedrock.Cfg.ResolvedProtocol() == aibconfig.BedrockProtocolInvokeModel
}

func (i *interceptionBase) Model() string {
Expand All @@ -164,7 +163,7 @@ func (i *interceptionBase) Model() string {
}

func (i *interceptionBase) baseTraceAttributes(r *http.Request, streaming bool) []attribute.KeyValue {
return []attribute.KeyValue{
attrs := []attribute.KeyValue{
attribute.String(tracing.RequestPath, r.URL.Path),
attribute.String(tracing.InterceptionID, i.id.String()),
attribute.String(tracing.InitiatorID, aibcontext.ActorIDFromContext(r.Context())),
Expand All @@ -173,6 +172,10 @@ func (i *interceptionBase) baseTraceAttributes(r *http.Request, streaming bool)
attribute.Bool(tracing.Streaming, streaming),
attribute.Bool(tracing.IsBedrock, i.bedrock != nil),
}
if i.bedrock != nil {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P3 [CRF-2] The new aws_bedrock_protocol trace branch, the PR's only observable behavior, has 0% coverage, and ResolvedProtocol()'s empty->invoke-model mapping never runs under test. (Bisky P3; Netero Note)

go tool cover reports both baseTraceAttributes (base.go:166) and ResolvedProtocol (config.go:76) at 0.0%

If the mapping regresses (emits "", or the branch stops firing), spans silently carry a missing or wrong protocol and nobody notices until someone debugs a Bedrock incident with bad telemetry. Blast radius is observability only, hence P3, but the gap is permanent absent a test. Bisky wrote and ran a green sketch: build an interceptionBase with a bedrock config, assert baseTraceAttributes emits invoke-model for the empty protocol and omits the attribute when bedrock == nil. One test covers the resolver, the append branch, and the nil-guard.

🤖

attrs = append(attrs, attribute.String(tracing.BedrockProtocol, string(i.bedrock.Cfg.ResolvedProtocol())))
}
return attrs
}

func (i *interceptionBase) injectTools() {
Expand Down
9 changes: 9 additions & 0 deletions aibridge/internal/integrationtest/trace_internal_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -187,6 +187,9 @@ func TestTraceAnthropic(t *testing.T) {
attribute.Bool(tracing.Streaming, tc.streaming),
attribute.Bool(tracing.IsBedrock, tc.bedrock),
}
if tc.bedrock {
attrs = append(attrs, attribute.String(tracing.BedrockProtocol, string(config.BedrockProtocolInvokeModel)))
}

require.Len(t, sr.Ended(), totalCount)
verifyTraces(t, sr, tc.expect, attrs)
Expand Down Expand Up @@ -311,6 +314,9 @@ func TestTraceAnthropicErr(t *testing.T) {
attribute.Bool(tracing.Streaming, tc.streaming),
attribute.Bool(tracing.IsBedrock, tc.bedrock),
}
if tc.bedrock {
attrs = append(attrs, attribute.String(tracing.BedrockProtocol, string(config.BedrockProtocolInvokeModel)))
}

verifyTraces(t, sr, tc.expect, attrs)
})
Expand Down Expand Up @@ -422,6 +428,9 @@ func TestInjectedToolsTrace(t *testing.T) {
}
if tc.expectProvider == config.ProviderAnthropic {
attrs = append(attrs, attribute.Bool(tracing.IsBedrock, tc.bedrock))
if tc.bedrock {
attrs = append(attrs, attribute.String(tracing.BedrockProtocol, string(config.BedrockProtocolInvokeModel)))
}
}

verifyTraces(t, sr, []expectTrace{{"Intercept.ProcessRequest.ToolCall", 1, codes.Unset}}, attrs)
Expand Down
3 changes: 3 additions & 0 deletions aibridge/provider/anthropic.go
Original file line number Diff line number Diff line change
Expand Up @@ -79,6 +79,9 @@ func NewAnthropic(ctx context.Context, cfg config.Anthropic, bedrockCfg *config.
if runtimeCfg.Region == "" {
runtimeCfg.Region = resolvedRegion
}
if err := runtimeCfg.Validate(); err != nil {
return nil, xerrors.Errorf("bedrock config: %w", err)
}
bedrock = &messages.BedrockRuntime{Cfg: runtimeCfg, Creds: creds}
}

Expand Down
36 changes: 36 additions & 0 deletions aibridge/provider/anthropic_internal_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -114,6 +114,42 @@ func TestNewAnthropic_KeyResolution(t *testing.T) {
}
}

// NOTE: no t.Parallel() because the subtests use t.Setenv.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

needs to be a nolint

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It passed CI

func TestNewAnthropic_BedrockRegionResolution(t *testing.T) {
t.Run("mantle_region_from_env", func(t *testing.T) {
t.Setenv("AWS_REGION", "us-west-2")

p, err := NewAnthropic(context.Background(), config.Anthropic{}, &config.AWSBedrock{
BaseURL: "https://bedrock-mantle.us-west-2.api.aws/anthropic",
Protocol: config.BedrockProtocolMantle,
AccessKey: "test-key",
AccessKeySecret: "test-secret",
})
require.NoError(t, err)
require.NotNil(t, p.bedrock)
require.Equal(t, "us-west-2", p.bedrock.Cfg.Region)
})

t.Run("mantle_no_region_anywhere", func(t *testing.T) {
// Clear every source the AWS SDK consults for a region so none
// resolves, then confirm construction rejects the mantle provider.
t.Setenv("AWS_REGION", "")
t.Setenv("AWS_DEFAULT_REGION", "")
t.Setenv("AWS_PROFILE", "")
t.Setenv("AWS_CONFIG_FILE", "/dev/null")
t.Setenv("AWS_SHARED_CREDENTIALS_FILE", "/dev/null")
t.Setenv("AWS_EC2_METADATA_DISABLED", "true")

_, err := NewAnthropic(context.Background(), config.Anthropic{}, &config.AWSBedrock{
BaseURL: "https://proxy.internal",
Protocol: config.BedrockProtocolMantle,
AccessKey: "test-key",
AccessKeySecret: "test-secret",
})
require.ErrorContains(t, err, "region required")
})
}

func TestAnthropic_CreateInterceptor(t *testing.T) {
t.Parallel()

Expand Down
13 changes: 7 additions & 6 deletions aibridge/tracing/tracing.go
Original file line number Diff line number Diff line change
Expand Up @@ -17,12 +17,13 @@ const (
// trace attribute key constants
RequestPath = "request_path"

InterceptionID = "interception_id"
InitiatorID = "user_id"
Provider = "provider"
Model = "model"
Streaming = "streaming"
IsBedrock = "aws_bedrock"
InterceptionID = "interception_id"
InitiatorID = "user_id"
Provider = "provider"
Model = "model"
Streaming = "streaming"
IsBedrock = "aws_bedrock"
BedrockProtocol = "aws_bedrock_protocol"

PassthroughURL = "passthrough_url"
PassthroughUpstreamURL = "passthrough_upstream_url"
Expand Down
4 changes: 4 additions & 0 deletions cli/aibridged_internal_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -175,6 +175,8 @@ func TestBuildProviders(t *testing.T) {
cfg.LegacyBedrock.Region = serpent.String("us-west-2")
cfg.LegacyBedrock.AccessKey = serpent.String("AKID")
cfg.LegacyBedrock.AccessKeySecret = serpent.String("secret")
cfg.LegacyBedrock.Model = serpent.String("anthropic.claude-3-5-sonnet-20241022-v2:0")
cfg.LegacyBedrock.SmallFastModel = serpent.String("anthropic.claude-3-5-haiku-20241022-v1:0")

providers, err := buildFromEnv(t, cfg)
require.NoError(t, err)
Expand All @@ -191,6 +193,8 @@ func TestBuildProviders(t *testing.T) {
cfg.LegacyBedrock.Region = serpent.String("us-west-2")
cfg.LegacyBedrock.AccessKey = serpent.String("AKID")
cfg.LegacyBedrock.AccessKeySecret = serpent.String("secret")
cfg.LegacyBedrock.Model = serpent.String("anthropic.claude-3-5-sonnet-20241022-v2:0")
cfg.LegacyBedrock.SmallFastModel = serpent.String("anthropic.claude-3-5-haiku-20241022-v1:0")

providers, err := buildFromEnv(t, cfg)
require.NoError(t, err)
Expand Down
2 changes: 2 additions & 0 deletions cli/server_aibridge_internal_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -727,6 +727,8 @@ func TestBuildProviderFromProtoSetsAPIDumpDir(t *testing.T) {
Region: "us-east-1",
AccessKey: "AKID",
AccessKeySecret: "secret",
Model: "anthropic.claude-3-5-sonnet-20241022-v2:0",
SmallFastModel: "anthropic.claude-3-5-haiku-20241022-v1:0",
},
},
expectedType: aibridge.ProviderAnthropic,
Expand Down
Loading