From ca3bde259b40bb1d2ecb1c29a82a01e41c73e8ef Mon Sep 17 00:00:00 2001 From: akshaydeo Date: Mon, 27 Jul 2026 09:33:08 +0530 Subject: [PATCH] test fixes --- core/providers/openai/responses_test.go | 24 +++++- core/schemas/responses.go | 26 +++++++ framework/sidekiq/sidekiq_test.go | 15 +++- .../handlers/mcpserver_auth_test.go | 25 +++++-- transports/bifrost-http/lib/config_test.go | 33 ++++++++ transports/config.schema.json | 75 ++++++++++++++++++- 6 files changed, 188 insertions(+), 10 deletions(-) diff --git a/core/providers/openai/responses_test.go b/core/providers/openai/responses_test.go index 5b318b88cc4..72ec57dd167 100644 --- a/core/providers/openai/responses_test.go +++ b/core/providers/openai/responses_test.go @@ -480,6 +480,10 @@ func TestToOpenAIResponsesRequest_GPTOSS_SummaryToContentBlocks(t *testing.T) { message schemas.ResponsesMessage expectedBlocks int expectedBlockText string + // OpenAI accepts `role` only on message input items, so ToOpenAIResponsesRequest + // strips it from a reasoning item that carries no explicit type. Set this for + // those cases: the role is expected to be dropped, not preserved. + expectRoleDropped bool description string }{ { @@ -533,6 +537,9 @@ func TestToOpenAIResponsesRequest_GPTOSS_SummaryToContentBlocks(t *testing.T) { }, expectedBlocks: 1, expectedBlockText: "existing content", + // Typeless reasoning item, so role is stripped here too — the strip + // happens before the summary-conversion branch is chosen. + expectRoleDropped: true, description: "gpt-oss model should preserve message when Content already exists", }, { @@ -552,6 +559,8 @@ func TestToOpenAIResponsesRequest_GPTOSS_SummaryToContentBlocks(t *testing.T) { }, expectedBlocks: 1, expectedBlockText: "variant summary", + // No explicit Type on a reasoning item, so role is stripped. + expectRoleDropped: true, description: "gpt-oss variant model should also convert Summary to ContentBlocks", }, } @@ -611,9 +620,6 @@ func TestToOpenAIResponsesRequest_GPTOSS_SummaryToContentBlocks(t *testing.T) { if tt.message.Status != nil && (resultMsg.Status == nil || *resultMsg.Status != *tt.message.Status) { t.Errorf("Expected Status to be preserved") } - if tt.message.Role != nil && (resultMsg.Role == nil || *resultMsg.Role != *tt.message.Role) { - t.Errorf("Expected Role to be preserved") - } } else { // For other cases, verify message is preserved as-is if resultMsg.Content != nil && len(resultMsg.Content.ContentBlocks) > 0 { @@ -622,6 +628,18 @@ func TestToOpenAIResponsesRequest_GPTOSS_SummaryToContentBlocks(t *testing.T) { } } } + + // Role handling is asserted for every case, not just the + // summary-conversion one: the converter strips role from any non-message + // item before it reaches either branch, so scoping this check to + // Content == nil would leave the existing-content path unverified. + if tt.expectRoleDropped { + if resultMsg.Role != nil { + t.Errorf("Expected Role to be dropped for a typeless reasoning item, got %q", *resultMsg.Role) + } + } else if tt.message.Role != nil && (resultMsg.Role == nil || *resultMsg.Role != *tt.message.Role) { + t.Errorf("Expected Role to be preserved") + } }) } } diff --git a/core/schemas/responses.go b/core/schemas/responses.go index c4d9321e452..5ec3bb53220 100644 --- a/core/schemas/responses.go +++ b/core/schemas/responses.go @@ -1283,6 +1283,23 @@ func (m *ResponsesMessage) UnmarshalJSON(data []byte) error { // consumers that read Arguments keep working; MarshalJSON still re-emits the // preserved bytes verbatim, so this is additive and does not affect round-trip. m.setToolArguments(json.RawMessage(gjson.GetBytes(data, "arguments").Raw)) + // Same rationale for `execution`: Codex reads it to decide whether to + // dispatch the call client-side, and returning early here skips the + // field-level decode that would otherwise populate it. + if execution := gjson.GetBytes(data, "execution"); execution.Type == gjson.String && execution.String() != "" { + if m.ResponsesToolMessage == nil { + m.ResponsesToolMessage = &ResponsesToolMessage{} + } + m.Execution = Ptr(execution.String()) + } + // tool_search_output carries the discovered tool list; surface it for the + // same reason. Held as raw JSON because the embedded ResponsesMCPListTools + // decode drops the per-tool type discriminator. + if t == string(ResponsesMessageTypeToolSearchOutput) { + if tools := gjson.GetBytes(data, "tools"); tools.IsArray() { + m.ToolSearchOutputTools = json.RawMessage(tools.Raw) + } + } m.rawPreserved = append([]byte(nil), data...) return nil } @@ -1350,6 +1367,15 @@ func responsesToolArgumentsToString(raw json.RawMessage) string { func (m ResponsesMessage) MarshalJSON() ([]byte, error) { type Alias ResponsesMessage + // Items decoded through the raw-preserved fast path are re-emitted byte for + // byte. That path deliberately skips field-level decoding, so the struct holds + // only Type plus the few fields surfaced for downstream readers — marshalling + // from those fields would silently drop everything else on the item (id, + // status, call_id, per-tool type discriminators), which OpenAI then rejects. + if len(m.rawPreserved) > 0 { + return append([]byte(nil), m.rawPreserved...), nil + } + // Re-emit the raw tools captured during unmarshal so the type discriminator survives. if m.Type != nil && *m.Type == ResponsesMessageTypeToolSearchOutput { aux := &struct { diff --git a/framework/sidekiq/sidekiq_test.go b/framework/sidekiq/sidekiq_test.go index 3e12cdb575f..83843e25adc 100644 --- a/framework/sidekiq/sidekiq_test.go +++ b/framework/sidekiq/sidekiq_test.go @@ -183,6 +183,14 @@ func testRunner(store Store) *Runner { return New(store, bifrost.NewDefaultLogger(schemas.LogLevelError), 4, "") } +// testRunnerWithID builds a runner that identifies itself. An empty runner ID +// makes New set staleAfter to 0, so every running job reads as instantly stale +// and is immediately reclaimable — fine for a lone runner recovering its own +// leftovers, wrong for any test where two runners contend. +func testRunnerWithID(store Store, runnerID string) *Runner { + return New(store, bifrost.NewDefaultLogger(schemas.LogLevelError), 4, runnerID) +} + func TestEnqueueRunsHandlerAndCompletes(t *testing.T) { store := newFakeStore() r := testRunner(store) @@ -312,7 +320,12 @@ func TestRunnerRaceSingleWinner(t *testing.T) { handler := func(_ context.Context, job tables.TableSidekiqJob, _ ProgressFunc) (string, error) { return "done", nil } - r1, r2 := testRunner(store), testRunner(store) + // Distinct runner IDs, as the doc comment above requires. With the empty ID + // testRunner uses, staleAfter is 0, so the loser reads the winner's + // microseconds-old claim as stale and re-claims it — the job then runs twice + // and the original owner cannot complete it. That is the documented behaviour + // of an unidentified runner, not a claim bug, but it makes this test racy. + r1, r2 := testRunnerWithID(store, "runner-A"), testRunnerWithID(store, "runner-B") r1.Register("k", handler) r2.Register("k", handler) diff --git a/transports/bifrost-http/handlers/mcpserver_auth_test.go b/transports/bifrost-http/handlers/mcpserver_auth_test.go index 224570a5909..fff8d62479a 100644 --- a/transports/bifrost-http/handlers/mcpserver_auth_test.go +++ b/transports/bifrost-http/handlers/mcpserver_auth_test.go @@ -413,19 +413,34 @@ func TestGetMCPServerForRequest_PreAuthenticatedUserPath(t *testing.T) { assert.Contains(t, err.Error(), "no MCP access grant") }) - t.Run("stamped user id with a header VK is rejected as conflicting", func(t *testing.T) { + // A stamped user id wins over any virtual key sent in the request headers: the + // request is scoped to the user's representative VK and the header VK is never + // honoured. Rejecting the pair as conflicting is deliberately NOT this + // function's job — that decision belongs upstream, in the SCIM inference + // middleware, which applies the operator's dual_credential_conflict_behavior + // before any identity is stamped (see getMCPServerForRequest). What matters + // here is that a header VK cannot escalate or redirect an already-authenticated + // user to a different key. + t.Run("stamped user id takes precedence over a header VK", func(t *testing.T) { cfg := newTestOAuth2Config(newStore(), configtables.MCPServerAuthModeBoth, true) h := newTestMCPHandler(cfg) h.identityResolver = &fakeResolver{userVKID: "vk-row-1"} - h.vkMCPServers[activeVK.Value.GetValue()] = server.NewMCPServer("vk", "v0") + userVKServer := server.NewMCPServer("vk", "v0") + h.vkMCPServers[activeVK.Value.GetValue()] = userVKServer + headerVKServer := server.NewMCPServer("header-vk", "v0") + h.vkMCPServers["sk-bf-header"] = headerVKServer ctx := &fasthttp.RequestCtx{} ctx.SetUserValue(schemas.BifrostContextKeyUserID, "user-1") ctx.Request.Header.Set(string(schemas.BifrostContextKeyVirtualKey), "sk-bf-header") - _, err := h.getMCPServerForRequest(ctx) - require.Error(t, err) - assert.Contains(t, err.Error(), "conflicting credentials") + res, err := h.getMCPServerForRequest(ctx) + require.NoError(t, err) + require.NotNil(t, res) + assert.Equal(t, userVKServer, res.mcpServer) + assert.NotEqual(t, headerVKServer, res.mcpServer) + assert.Nil(t, res.jwtVK) + assert.Nil(t, res.jwtClaims) }) t.Run("inactive representative virtual key is rejected", func(t *testing.T) { diff --git a/transports/bifrost-http/lib/config_test.go b/transports/bifrost-http/lib/config_test.go index f910ff1e081..81ae99b334d 100644 --- a/transports/bifrost-http/lib/config_test.go +++ b/transports/bifrost-http/lib/config_test.go @@ -767,6 +767,39 @@ func (m *MockConfigStore) UpdateBudgets(ctx context.Context, budgets []*tables.T return nil } +// UpdateBudgetOverride mirrors RDBConfigStore: it applies only the override +// columns to the stored budget and returns the updated row, leaving usage and +// base configuration untouched. Reusing SetOverrideAt keeps the anchoring and +// validation identical to the real store rather than re-deriving it here, and +// an unknown id yields configstore.ErrNotFound as the RDB store does. +func (m *MockConfigStore) UpdateBudgetOverride(ctx context.Context, id string, amount float64, mode tables.BudgetOverrideMode, cyclesTotal int, calendarAligned bool, tx ...*gorm.DB) (*tables.TableBudget, error) { + if m.governanceConfig == nil { + return nil, configstore.ErrNotFound + } + for i := range m.governanceConfig.Budgets { + if m.governanceConfig.Budgets[i].ID != id { + continue + } + // Validate against a copy and commit only on success, mirroring + // RDBConfigStore.UpdateBudgetOverride: it loads the row into a local struct + // and returns before its Updates() call, so a rejected override persists + // nothing. Mutating the stored budget in place would leak IsCalendarAligned + // on failure — SetOverrideAt rolls back the override columns, not that flag. + // + // IsCalendarAligned is not persisted on the budget row, so the caller + // supplies it — same contract as the RDB store. + candidate := m.governanceConfig.Budgets[i] + candidate.IsCalendarAligned = calendarAligned + if err := candidate.SetOverrideAt(amount, mode, cyclesTotal, candidate.WindowStart(time.Now())); err != nil { + return nil, err + } + m.governanceConfig.Budgets[i] = candidate + updated := candidate + return &updated, nil + } + return nil, configstore.ErrNotFound +} + func (m *MockConfigStore) GetBudget(ctx context.Context, id string, tx ...*gorm.DB) (*tables.TableBudget, error) { return nil, nil } diff --git a/transports/config.schema.json b/transports/config.schema.json index 425140c6e90..d257373ee01 100644 --- a/transports/config.schema.json +++ b/transports/config.schema.json @@ -513,10 +513,83 @@ "type": "boolean", "description": "Deprecated: set calendar_aligned on the owner (team / virtual key / access profile) instead. Kept for backward compatibility with older config.json files; ignored unless a reconciler promotes it to its owner.", "default": false + }, + "override_amount": { + "type": "number", + "description": "Amount added to max_limit while an override is active", + "default": 0 + }, + "override_mode": { + "type": "string", + "enum": ["", "cycles", "forever"], + "description": "How long the override stays active: 'cycles' for a finite number of reset windows, 'forever' until explicitly removed, '' for no override", + "default": "" + }, + "override_cycles_remaining": { + "type": "integer", + "description": "Reset windows the current override still covers, including the current one. Derived state: recomputed from override_anchor_reset, override_cycles_total and last_reset, so it is never the source of truth.", + "default": 0 + }, + "override_cycles_total": { + "type": "integer", + "description": "Number of reset windows the current finite override was granted for. Immutable for the life of a grant.", + "default": 0 + }, + "override_anchor_reset": { + "type": "string", + "format": "date-time", + "description": "Window boundary at which the current finite override was granted. Immutable for the life of a grant; absent when there is no finite override." } }, "required": ["id", "max_limit", "reset_duration"], - "additionalProperties": false + "additionalProperties": false, + "allOf": [ + { + "if": { + "not": { + "properties": { "override_mode": { "enum": ["cycles", "forever"] } }, + "required": ["override_mode"] + } + }, + "then": { + "properties": { + "override_amount": { "const": 0 }, + "override_cycles_remaining": { "const": 0 }, + "override_cycles_total": { "const": 0 } + }, + "not": { "required": ["override_anchor_reset"] } + } + }, + { + "if": { + "properties": { "override_mode": { "const": "cycles" } }, + "required": ["override_mode"] + }, + "then": { + "required": ["override_amount", "override_cycles_remaining", "override_cycles_total", "override_anchor_reset"], + "properties": { + "override_amount": { "exclusiveMinimum": 0 }, + "override_cycles_remaining": { "minimum": 1 }, + "override_cycles_total": { "minimum": 1 } + } + } + }, + { + "if": { + "properties": { "override_mode": { "const": "forever" } }, + "required": ["override_mode"] + }, + "then": { + "required": ["override_amount"], + "properties": { + "override_amount": { "exclusiveMinimum": 0 }, + "override_cycles_remaining": { "const": 0 }, + "override_cycles_total": { "const": 0 } + }, + "not": { "required": ["override_anchor_reset"] } + } + } + ] } }, "rate_limits": {