From 867c6205724cd072a66501028c7fe6abfbfe0f9d Mon Sep 17 00:00:00 2001 From: Jakub Hrozek Date: Wed, 23 Sep 2026 10:43:53 +0200 Subject: [PATCH 1/3] fix(enrollment): preserve workspace enrollment authority Workspace enrollment replaced the session's whole tool capability set with the latest broker catalogue, discarding unrelated authority. Restart and successor paths had no durable way to tell stale broker keys from deliberately retained capabilities, and a cross-source registration collision could make a durable key resolve to a different implementation. Persist a bounded broker-key ledger with a presence bit on the Session aggregate, through sessnap and eventsource. The server sets authority and ledger together on every root-creation and authority-carrying path (clear, fork, and CreateSession with a source session); session.New starts with absent provenance, so a path that forgets the ledger refuses enrollment. Completion replaces only the prior broker bundle and keeps excluded keys in the ledger, so an exclusion survives a refresh that omits the key. Composition registers broker tools last against frozen registration keys and reports a collision as a typed, presentation-safe error that unwraps to failed precondition. A recorded broker key still held in authority and now owned by a non-broker tool fails every engine build. An enrollment-time collision or ledger overflow clears the pending enrollment; mecatui shows it as a terminal failure without sending a retained prompt. Offline tests cover every acceptance criterion in docs/acceptance/workspace-enrollment-authority.md: aggregate, persistence, restart, successors, composition (the real Build with a configured deny), server end to end, and TUI. Implements the plan in #1741. Relates to #1740. Co-Authored-By: mecatl --- cmd/mecatui/ui/workspace_enrollment.go | 13 + cmd/mecatui/ui/workspace_enrollment_test.go | 85 ++- engine/CHANGELOG.md | 13 + engine/adapter/eventsource/eventsource.go | 7 + .../workspace_enrollment_authority_test.go | 58 ++ engine/adapter/sessnap/sessnap.go | 87 +++ .../workspace_enrollment_authority_test.go | 134 ++++ engine/api/session.txt | 2 + engine/session/session.go | 6 + engine/session/workspace_enrollment.go | 107 ++- .../workspace_enrollment_authority_test.go | 197 ++++++ engine/session/workspace_enrollment_test.go | 3 + .../mcpbroker/query_continuation_test.go | 6 +- .../server/authorization_close_test.go | 2 +- internal/adapter/server/broker_rebind_test.go | 2 +- .../adapter/server/harness_context_test.go | 6 +- internal/adapter/server/mcp_authorization.go | 2 +- .../mcp_authorization_lifecycle_test.go | 6 +- .../mcp_authorization_transport_test.go | 2 +- internal/adapter/server/mcp_broker.go | 7 +- .../server/mcp_broker_id_validation_test.go | 2 +- .../mcp_broker_multi_upstream_e2e_test.go | 18 +- internal/adapter/server/mcp_broker_test.go | 37 +- .../adapter/server/placement_successor.go | 5 +- .../server/placement_successor_broker_test.go | 100 ++- internal/adapter/server/service.go | 224 +++++-- .../adapter/server/workspace_enrollment.go | 60 +- .../server/workspace_enrollment_test.go | 617 +++++++++++++++++- internal/app/build.go | 93 ++- internal/app/catalog.go | 120 +++- internal/app/catalog_classification.go | 20 +- internal/app/catalog_classification_test.go | 87 +++ .../harness_context_execution_source_test.go | 2 +- internal/app/mcp_broker_p10_test.go | 53 +- internal/app/mcp_runtime.go | 4 +- .../mcpbroker/workspace_enrollment_test.go | 2 +- .../mecak8s/identity-and-client-access.md | 8 + 37 files changed, 1969 insertions(+), 228 deletions(-) create mode 100644 engine/adapter/eventsource/workspace_enrollment_authority_test.go create mode 100644 engine/adapter/sessnap/workspace_enrollment_authority_test.go create mode 100644 engine/session/workspace_enrollment_authority_test.go diff --git a/cmd/mecatui/ui/workspace_enrollment.go b/cmd/mecatui/ui/workspace_enrollment.go index e22bd4f1f0..618ab6abfc 100644 --- a/cmd/mecatui/ui/workspace_enrollment.go +++ b/cmd/mecatui/ui/workspace_enrollment.go @@ -250,6 +250,12 @@ func (m Model) applyWorkspaceEnrollment(msg workspaceEnrollmentMsg) (tea.Model, // deadline outcome. Leave it silent and prevent a queued poll re-arm. m.enrollment.presentationDelivered = false return m, nil + case isWorkspaceEnrollmentCollision(msg.err.Error()): + _, notice, statusMsg := terminalWorkspaceEnrollmentNotice(client.WorkspaceEnrollmentFailed) + m.enrollment = workspaceEnrollmentState{Status: client.WorkspaceEnrollmentFailed, controlGen: m.enrollment.controlGen} + m.workspaceEnrollmentNotice = notice + m.statusMsg = statusMsg + return m, nil default: m.enrollment.err = oneLine(terminaltext.Sanitize(msg.err.Error())) m.workspaceEnrollmentNotice = m.enrollment.err @@ -360,6 +366,13 @@ func friendlyWorkspaceEnrollmentRejection(raw string) string { return raw } +func isWorkspaceEnrollmentCollision(raw string) bool { + return (strings.Contains(raw, "workspace tool ") && + strings.Contains(raw, "conflicts with an existing registration") || + strings.Contains(raw, "workspace enrollment registration collision")) && + strings.Contains(raw, "correct the broker configuration and retry") +} + func isWorkspaceEnrollmentRejection(raw string) bool { return strings.Contains(raw, "workspace services must be connected before prompting") } diff --git a/cmd/mecatui/ui/workspace_enrollment_test.go b/cmd/mecatui/ui/workspace_enrollment_test.go index 49ff903fe5..3b0b06b507 100644 --- a/cmd/mecatui/ui/workspace_enrollment_test.go +++ b/cmd/mecatui/ui/workspace_enrollment_test.go @@ -8,6 +8,8 @@ import ( "testing" tea "charm.land/bubbletea/v2" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" "github.com/stacklok/mecatl/cmd/mecatui/client" ) @@ -344,6 +346,86 @@ func TestWorkspaceEnrollmentFailedDoesNotPollOrRetry(t *testing.T) { if control.connectCalls != 0 { t.Fatalf("failed enrollment rechecked %d times", control.connectCalls) } + terminal, _, guidance := terminalWorkspaceEnrollmentNotice(client.WorkspaceEnrollmentFailed) + if !terminal || !strings.Contains(guidance, "failed") || !strings.Contains(guidance, "/tools-connect") || strings.Contains(guidance, "bundle-1") { + t.Fatalf("terminal conflict copy = terminal:%t %q; want safe retry guidance", terminal, guidance) + } +} + +func TestWorkspaceEnrollmentAuthority_Scenario3_TerminalConflictUX(t *testing.T) { + collision := status.Error(codes.FailedPrecondition, `workspace tool "mcp__prior__tool" conflicts with an existing registration; correct the broker configuration and retry`) + control := &workspaceEnrollmentControlFake{connectErr: collision} + m, send := builtinDispatchModel(t, client.Capabilities{WorkspaceEnrollment: true}, false) + m.deps.WorkspaceEnrollment = control + m.pendingInitialPrompt = "list my open pull requests" + + mm, connect := m.runToolsConnect() + m = mm.(Model) + if connect == nil { + t.Fatal("/tools-connect returned no command") + } + m = applyAll(m, connect()) + if m.enrollment.Status != client.WorkspaceEnrollmentFailed || !strings.Contains(m.statusMsg, "failed") || !strings.Contains(m.statusMsg, "/tools-connect") { + t.Fatalf("collision did not show terminal failure and retry: enrollment=%+v status=%q", m.enrollment, m.statusMsg) + } + if m.pendingInitialPrompt != "list my open pull requests" || len(promptTexts(send)) != 0 { + t.Fatalf("terminal collision sent retained prompt: pending=%q sent=%v", m.pendingInitialPrompt, promptTexts(send)) + } + control.connect = client.WorkspaceEnrollment{ID: "bundle-retry", Status: client.WorkspaceEnrollmentConnected} + control.connectErr = nil + mm, retry := m.runToolsConnect() + m = mm.(Model) + if retry == nil { + t.Fatal("terminal collision did not offer explicit /tools-connect retry") + } + connectedMsg := retry() + mm, connectedCmd := m.Update(connectedMsg) + m = mm.(Model) + if connectedCmd == nil { + t.Fatal("successful retry did not produce connected state updates") + } + outer, ok := connectedCmd().(tea.BatchMsg) + if !ok || len(outer) < 2 { + t.Fatalf("connected retry command = %#v, want enrollment and prompt batches", outer) + } + firstBatchLeaf(t, outer[1]) + if m.cancelRun != nil { + m.cancelRun() + } + if m.enrollment.Status != client.WorkspaceEnrollmentConnected || m.pendingInitialPrompt != "" { + t.Fatalf("successful retry did not finalize enrollment and consume the retained prompt: enrollment=%+v pending=%q", m.enrollment, m.pendingInitialPrompt) + } + prompts := promptTexts(send) + if len(prompts) != 1 || prompts[0] != "list my open pull requests" { + t.Fatalf("successful retry sent prompts %q, want the retained prompt exactly once", prompts) + } + if control.connectCalls != 2 { + t.Fatalf("connect calls = %d, want initial attempt and explicit retry", control.connectCalls) + } +} + +func TestWorkspaceEnrollmentAuthority_EmptyKeyConflictUsesSafeRetryUX(t *testing.T) { + collision := status.Error(codes.FailedPrecondition, "workspace enrollment registration collision; correct the broker configuration and retry") + control := &workspaceEnrollmentControlFake{connectErr: collision} + m, _ := builtinDispatchModel(t, client.Capabilities{WorkspaceEnrollment: true}, false) + m.deps.WorkspaceEnrollment = control + + mm, connect := m.runToolsConnect() + m = mm.(Model) + if connect == nil { + t.Fatal("/tools-connect returned no command") + } + m = applyAll(m, connect()) + + if m.enrollment.Status != client.WorkspaceEnrollmentFailed { + t.Fatalf("empty-key collision enrollment = %+v, want failed", m.enrollment) + } + if !strings.Contains(m.statusMsg, "failed") || !strings.Contains(m.statusMsg, "/tools-connect") { + t.Fatalf("empty-key collision did not show safe retry guidance: %q", m.statusMsg) + } + if strings.Contains(m.statusMsg, "registration collision") { + t.Fatalf("empty-key collision leaked raw RPC detail: %q", m.statusMsg) + } } // TestWorkspaceEnrollmentExpiredResetsInsteadOfPollingForever pins the bug @@ -446,13 +528,14 @@ func TestWorkspaceEnrollmentQueuedPresentationIsCancelledBeforeOpen(t *testing.T type workspaceEnrollmentControlFake struct { connect client.WorkspaceEnrollment + connectErr error connectCalls int cancelCalls int } func (f *workspaceEnrollmentControlFake) ConnectWorkspaceServices(context.Context, string) (client.WorkspaceEnrollment, error) { f.connectCalls++ - return f.connect, nil + return f.connect, f.connectErr } func (*workspaceEnrollmentControlFake) RetryWorkspaceEnrollment(context.Context, string, string) (client.WorkspaceEnrollment, error) { diff --git a/engine/CHANGELOG.md b/engine/CHANGELOG.md index 2e8fd643b5..60441cfa81 100644 --- a/engine/CHANGELOG.md +++ b/engine/CHANGELOG.md @@ -167,6 +167,12 @@ The covered surface is the eight core packages (`session`, `governance`, `learni execute, learn permission rules, arm hook waivers, or transition plan mode. Changed (breaking, pre-v1 minor). +- **Durable workspace-enrollment broker authority provenance** — adds + `session.Session.WorkspaceEnrollmentBrokerKeys` and + `session.Session.RestoreWorkspaceEnrollmentBrokerKeys` for exact bounded broker + registration-key provenance across snapshot/event-source restoration and + authority-carrying successors. Restore is write-once per aggregate. Added (minor). + - **Request-manifest schema-byte evidence** — adds `session.RequestManifestPayload.AdvertisedToolSchemaBytes` and exposes it through the target-bound debugger manifest view. Adds catalog registration-key metadata accessors so manifest enumeration does not refresh live tool specifications. Added (minor). - **Atomic ordinary permission-ask resolution** — adds `agent.AskResolution`, @@ -506,6 +512,13 @@ The covered surface is the eight core packages (`session`, `governance`, `learni decoding validates canonical fields, ignores unknown fields, and treats an omitted empty token ledger as its zero value. Changed/breaking (pre-v1 minor). +- **Workspace-enrollment authority replacement** — + `session.New` starts with legacy-absent broker provenance, and + `session.Session.CompleteWorkspaceEnrollment` now replaces only the prior + recorded broker bundle while retaining excluded broker keys in its ledger + across omitted refreshes, and refuses legacy-absent provenance. Changed (breaking, + pre-v1 minor). + - **Go compatibility floor** - the engine module requires Go 1.27. The root, provider, and authentication modules use the same floor. Changed (breaking, pre-v1 minor). diff --git a/engine/adapter/eventsource/eventsource.go b/engine/adapter/eventsource/eventsource.go index 520f9a8b1f..3d48e62743 100644 --- a/engine/adapter/eventsource/eventsource.go +++ b/engine/adapter/eventsource/eventsource.go @@ -136,6 +136,10 @@ type SessionMeta struct { // metadata. Nil is a documented legacy record written without authority; a // present payload is validated and bound before reconstruction proceeds. Authority *session.Authority + // WorkspaceEnrollmentBrokerKeys and its presence bit must come from the same + // committed authority revision. They are not reconstructed from events. + WorkspaceEnrollmentBrokerKeys []string + WorkspaceEnrollmentBrokerKeysPresent bool // ExternalBinding is the opaque composition-issued process-external // identity (e.g. an MCP broker attachment binding). Not event-carried: // safe to omit for a pure fold UNLESS the host requires exact external- @@ -199,6 +203,9 @@ func Fold(meta SessionMeta, events iter.Seq2[session.Event, error]) (*session.Se if err := restoreAuthority(s, meta.Authority); err != nil { return nil, fmt.Errorf("%w: %w", ErrReconstruct, err) } + if err := s.RestoreWorkspaceEnrollmentBrokerKeys(meta.WorkspaceEnrollmentBrokerKeys, meta.WorkspaceEnrollmentBrokerKeysPresent); err != nil { + return nil, fmt.Errorf("%w: restore workspace enrollment broker keys: %w", ErrReconstruct, err) + } if err := s.RestoreLabels(meta.Owner, session.Authority{}); err != nil { return nil, fmt.Errorf("%w: %w", ErrReconstruct, err) } diff --git a/engine/adapter/eventsource/workspace_enrollment_authority_test.go b/engine/adapter/eventsource/workspace_enrollment_authority_test.go new file mode 100644 index 0000000000..b3ae708783 --- /dev/null +++ b/engine/adapter/eventsource/workspace_enrollment_authority_test.go @@ -0,0 +1,58 @@ +package eventsource_test + +import ( + "reflect" + "testing" + + "github.com/stacklok/mecatl/engine/adapter/eventsource" + "github.com/stacklok/mecatl/engine/governance" + "github.com/stacklok/mecatl/engine/session" +) + +func TestWorkspaceEnrollmentAuthority_Scenario2_EventSourceRestoresProvenanceBeforeAwaiting(t *testing.T) { + authority := session.Authority{CapabilitySet: governance.CapabilitySet{Tools: []string{"Read", "broker-old"}}, Provenance: "test"} + call := toolCall("c1", "Shell", `{"command":"ls"}`) + ask := session.PendingAsk{AskID: "s1:0:c1:r0", Tool: "Shell", Reason: "needs approval"} + for _, tc := range []struct { + name string + keys []string + present bool + }{ + {name: "absent"}, + {name: "present empty", present: true}, + {name: "present nonempty", keys: []string{"broker-old"}, present: true}, + } { + t.Run(tc.name, func(t *testing.T) { + for _, awaiting := range []bool{false, true} { + t.Run(map[bool]string{false: "ordinary", true: "awaiting"}[awaiting], func(t *testing.T) { + m := meta() + m.Authority = &authority + m.WorkspaceEnrollmentBrokerKeys = tc.keys + m.WorkspaceEnrollmentBrokerKeysPresent = tc.present + events := []session.Event(nil) + if awaiting { + events = []session.Event{{Type: session.EvTurnStart, Turn: 0}, {Type: session.EvToolCall, Turn: 0, ToolCall: &call}, {Type: session.EvPermissionAsk, Turn: 0, Ask: &ask}} + } + s, err := eventsource.Fold(m, seq(events)) + if err != nil { + t.Fatalf("Fold: %v", err) + } + if awaiting && s.State != session.StateAwaiting { + t.Fatalf("state = %q, want awaiting", s.State) + } + keys, present := s.WorkspaceEnrollmentBrokerKeys() + if present != tc.present || !reflect.DeepEqual(keys, tc.keys) { + t.Fatalf("ledger = %q, %t; want %q, %t", keys, present, tc.keys, tc.present) + } + }) + } + }) + } + + m := meta() + m.Authority = &authority + m.WorkspaceEnrollmentBrokerKeys = []string{"broker-old"} + if _, err := eventsource.Fold(m, seq(nil)); err == nil { + t.Fatal("Fold accepted invalid metadata") + } +} diff --git a/engine/adapter/sessnap/sessnap.go b/engine/adapter/sessnap/sessnap.go index 50889325c5..b3d38c4d49 100644 --- a/engine/adapter/sessnap/sessnap.go +++ b/engine/adapter/sessnap/sessnap.go @@ -28,7 +28,9 @@ import ( "encoding/json" "errors" "fmt" + "strconv" "time" + "unicode/utf8" "github.com/stacklok/mecatl/engine/session" ) @@ -147,6 +149,10 @@ type Snapshot struct { // one governance.CapabilitySet representation or restore fails closed. Authority *session.Authority `json:"authority,omitempty"` ExecutionCapabilities *session.ExecutionCapabilities `json:"execution_capabilities,omitempty"` + // WorkspaceEnrollmentBrokerKeys is the exact completed broker registration-key + // ledger. Presence distinguishes known empty provenance from legacy absence. + WorkspaceEnrollmentBrokerKeys []string `json:"workspace_enrollment_broker_keys,omitempty"` + WorkspaceEnrollmentBrokerKeysPresent bool `json:"workspace_enrollment_broker_keys_present,omitempty"` // EnvironmentRef is the sole durable execution-environment identity. It is // required and must contain the exact provider revision used for reattachment. EnvironmentRef session.EnvironmentRef `json:"environment_ref"` @@ -330,6 +336,7 @@ func Of(s *session.Session) (Snapshot, error) { if occupancy, ok := s.LatestContextOccupancy(); ok { snap.LatestContextOccupancy = &occupancy } + snap.WorkspaceEnrollmentBrokerKeys, snap.WorkspaceEnrollmentBrokerKeysPresent = s.WorkspaceEnrollmentBrokerKeys() if s.Conversation != nil { snap.Messages = make([]messageDTO, len(s.Conversation.Messages)) for i, m := range s.Conversation.Messages { @@ -376,6 +383,9 @@ func (s Snapshot) Restore() (*session.Session, error) { if err := restoreAuthority(restored, s.Authority); err != nil { return nil, err } + if err := restored.RestoreWorkspaceEnrollmentBrokerKeys(s.WorkspaceEnrollmentBrokerKeys, s.WorkspaceEnrollmentBrokerKeysPresent); err != nil { + return nil, fmt.Errorf("sessnap: restore workspace enrollment broker keys: %w", err) + } // Rebuild the conversation history verbatim. for _, dto := range s.Messages { @@ -607,6 +617,13 @@ func Unmarshal(line []byte) (*session.Session, error) { if err := validateAuthorityWireClaim(wire.Authority); err != nil { return nil, fmt.Errorf("sessnap: decode authority: %w", err) } + var fields map[string]json.RawMessage + if err := json.Unmarshal(line, &fields); err != nil { + return nil, fmt.Errorf("sessnap: decode snapshot: %w", err) + } + if err := validateWorkspaceEnrollmentBrokerKeysWire(fields); err != nil { + return nil, fmt.Errorf("sessnap: decode workspace enrollment broker keys: %w", err) + } var snap Snapshot if err := json.Unmarshal(line, &snap); err != nil { @@ -615,6 +632,76 @@ func Unmarshal(line []byte) (*session.Session, error) { return snap.Restore() } +func validWorkspaceEnrollmentKeyEscapes(raw []byte) bool { + for i := 0; i < len(raw); i++ { + if raw[i] != '\\' { + continue + } + if i+1 >= len(raw) { + return false + } + if raw[i+1] != 'u' { + i++ + continue + } + if i+6 > len(raw) { + return false + } + code, ok := jsonEscapeCode(raw[i+2 : i+6]) + if !ok { + return false + } + i += 5 + if code < 0xd800 || code > 0xdfff { + continue + } + if code >= 0xdc00 || i+6 >= len(raw) || raw[i+1] != '\\' || raw[i+2] != 'u' { + return false + } + low, ok := jsonEscapeCode(raw[i+3 : i+7]) + if !ok || low < 0xdc00 || low > 0xdfff { + return false + } + i += 6 + } + return true +} + +func jsonEscapeCode(raw []byte) (rune, bool) { + code, err := strconv.ParseUint(string(raw), 16, 16) + if err != nil { + return 0, false + } + return rune(code), true +} + +func validateWorkspaceEnrollmentBrokerKeysWire(fields map[string]json.RawMessage) error { + keysRaw, keysSupplied := fields["workspace_enrollment_broker_keys"] + presentRaw, presentSupplied := fields["workspace_enrollment_broker_keys_present"] + present := false + if presentSupplied { + if bytes.Equal(bytes.TrimSpace(presentRaw), []byte("null")) { + return errors.New("null presence") + } + if err := json.Unmarshal(presentRaw, &present); err != nil { + return errors.New("presence must be a boolean") + } + } + var keys []string + if keysSupplied && !bytes.Equal(bytes.TrimSpace(keysRaw), []byte("null")) { + if !utf8.Valid(keysRaw) || !validWorkspaceEnrollmentKeyEscapes(keysRaw) { + return errors.New("keys contain invalid UTF-8") + } + if err := json.Unmarshal(keysRaw, &keys); err != nil { + return errors.New("keys must be an array of strings") + } + } + if (!present && len(keys) != 0) || !session.ValidWorkspaceEnrollmentToolNames(keys) { + return errors.New("invalid broker-key ledger") + } + return nil +} + func validateAuthorityWireClaim(raw json.RawMessage) error { if len(raw) == 0 { return nil // Genuine legacy record: no authority field. diff --git a/engine/adapter/sessnap/workspace_enrollment_authority_test.go b/engine/adapter/sessnap/workspace_enrollment_authority_test.go new file mode 100644 index 0000000000..8ab6f2cf51 --- /dev/null +++ b/engine/adapter/sessnap/workspace_enrollment_authority_test.go @@ -0,0 +1,134 @@ +package sessnap + +import ( + "encoding/json" + "fmt" + "reflect" + "strings" + "testing" +) + +func TestWorkspaceEnrollmentAuthority_Scenario2_ProvenancePresenceAndBounds(t *testing.T) { + s, pending := enrollmentSnapshotSession(t) + if err := s.AbortWorkspaceEnrollment(pending.ID); err != nil { + t.Fatalf("AbortWorkspaceEnrollment: %v", err) + } + if err := s.RestoreWorkspaceEnrollmentBrokerKeys([]string{"broker-old"}, true); err != nil { + t.Fatalf("RestoreWorkspaceEnrollmentBrokerKeys: %v", err) + } + encoded, err := Marshal(s) + if err != nil { + t.Fatalf("Marshal: %v", err) + } + var fields map[string]json.RawMessage + if err := json.Unmarshal(encoded, &fields); err != nil { + t.Fatal(err) + } + if string(fields["workspace_enrollment_broker_keys_present"]) != "true" { + t.Fatalf("presence JSON = %s, want true", fields["workspace_enrollment_broker_keys_present"]) + } + restored, err := Unmarshal(encoded) + if err != nil { + t.Fatalf("Unmarshal: %v", err) + } + keys, present := restored.WorkspaceEnrollmentBrokerKeys() + if !present || !reflect.DeepEqual(keys, []string{"broker-old"}) { + t.Fatalf("restored ledger = %q, %t", keys, present) + } + + base := `{"id":"legacy","state":"idle","mode":"default","limits":{},"counters":{},"environment_ref":{"Kind":"local","ID":".","Revision":"in-tree-v1"},"created_at":"1970-01-01T00:00:01Z","authority":{"capability_set":{"tools":["Read"]},"provenance":"test"}` + for name, suffix := range map[string]string{ + "legacy absent": `}`, + "present null": `,"workspace_enrollment_broker_keys":null,"workspace_enrollment_broker_keys_present":true}`, + "keys without presence": `,"workspace_enrollment_broker_keys":["broker"]}`, + "null presence": `,"workspace_enrollment_broker_keys_present":null}`, + "non-array keys": `,"workspace_enrollment_broker_keys":"broker","workspace_enrollment_broker_keys_present":true}`, + "duplicate keys": `,"workspace_enrollment_broker_keys":["broker","broker"],"workspace_enrollment_broker_keys_present":true}`, + "lone surrogate": `,"workspace_enrollment_broker_keys":["\ud800"],"workspace_enrollment_broker_keys_present":true}`, + } { + t.Run(name, func(t *testing.T) { + restored, err := Unmarshal([]byte(base + suffix)) + if name == "legacy absent" || name == "present null" { + if err != nil { + t.Fatalf("Unmarshal: %v", err) + } + keys, present := restored.WorkspaceEnrollmentBrokerKeys() + wantPresent := name == "present null" + if present != wantPresent || len(keys) != 0 { + t.Fatalf("ledger = %q, %t; want empty, %t", keys, present, wantPresent) + } + return + } + if err == nil { + t.Fatal("Unmarshal accepted invalid broker ledger JSON") + } + }) + } + + // Raw wire validation runs before JSON decoding can replace malformed bytes or + // lone escapes. Exercise every presence state and every documented boundary. + valid256 := make([]string, 256) + for i := range valid256 { + valid256[i] = fmt.Sprintf("k%03d", i) + } + tooMany := append(append([]string(nil), valid256...), "extra") + tooLong := []string{strings.Repeat("x", 257)} + // 128 distinct 256-byte keys are exactly 32 KiB; append one byte to exceed it. + exactTotal := make([]string, 128) + for i := range exactTotal { + exactTotal[i] = strings.Repeat("x", 252) + fmt.Sprintf("%04d", i) + } + overTotal := append(append([]string(nil), exactTotal...), "z") + encodeKeys := func(keys []string) string { + encoded, err := json.Marshal(keys) + if err != nil { + t.Fatal(err) + } + return string(encoded) + } + for _, tc := range []struct { + name, suffix string + wantPresent bool + wantKeys []string + valid bool + }{ + {"missing", `}`, false, nil, true}, + {"false omitted", `,"workspace_enrollment_broker_keys_present":false}`, false, nil, true}, + {"false empty", `,"workspace_enrollment_broker_keys":[],"workspace_enrollment_broker_keys_present":false}`, false, nil, true}, + {"false null", `,"workspace_enrollment_broker_keys":null,"workspace_enrollment_broker_keys_present":false}`, false, nil, true}, + {"true omitted", `,"workspace_enrollment_broker_keys_present":true}`, true, nil, true}, + {"true empty", `,"workspace_enrollment_broker_keys":[],"workspace_enrollment_broker_keys_present":true}`, true, nil, true}, + {"true null", `,"workspace_enrollment_broker_keys":null,"workspace_enrollment_broker_keys_present":true}`, true, nil, true}, + {"true keys", `,"workspace_enrollment_broker_keys":["broker"],"workspace_enrollment_broker_keys_present":true}`, true, []string{"broker"}, true}, + {"null presence", `,"workspace_enrollment_broker_keys_present":null}`, false, nil, false}, + {"nonboolean presence", `,"workspace_enrollment_broker_keys_present":"true"}`, false, nil, false}, + {"nonarray keys", `,"workspace_enrollment_broker_keys":"broker","workspace_enrollment_broker_keys_present":true}`, false, nil, false}, + {"nonstring key", `,"workspace_enrollment_broker_keys":[1],"workspace_enrollment_broker_keys_present":true}`, false, nil, false}, + {"empty key", `,"workspace_enrollment_broker_keys":[""],"workspace_enrollment_broker_keys_present":true}`, false, nil, false}, + {"control key", `,"workspace_enrollment_broker_keys":["bad\u0000"],"workspace_enrollment_broker_keys_present":true}`, false, nil, false}, + {"raw invalid utf8", ",\"workspace_enrollment_broker_keys\":[\"\xff\"],\"workspace_enrollment_broker_keys_present\":true}", false, nil, false}, + {"256 keys", `,"workspace_enrollment_broker_keys":` + encodeKeys(valid256) + `,"workspace_enrollment_broker_keys_present":true}`, true, valid256, true}, + {"257 keys", `,"workspace_enrollment_broker_keys":` + encodeKeys(tooMany) + `,"workspace_enrollment_broker_keys_present":true}`, false, nil, false}, + {"256 bytes", `,"workspace_enrollment_broker_keys":` + encodeKeys([]string{strings.Repeat("x", 256)}) + `,"workspace_enrollment_broker_keys_present":true}`, true, []string{strings.Repeat("x", 256)}, true}, + {"257 bytes", `,"workspace_enrollment_broker_keys":` + encodeKeys(tooLong) + `,"workspace_enrollment_broker_keys_present":true}`, false, nil, false}, + {"32 KiB", `,"workspace_enrollment_broker_keys":` + encodeKeys(exactTotal) + `,"workspace_enrollment_broker_keys_present":true}`, true, exactTotal, true}, + {"over 32 KiB", `,"workspace_enrollment_broker_keys":` + encodeKeys(overTotal) + `,"workspace_enrollment_broker_keys_present":true}`, false, nil, false}, + } { + t.Run(tc.name, func(t *testing.T) { + restored, err := Unmarshal([]byte(base + tc.suffix)) + if !tc.valid { + if err == nil { + t.Fatal("Unmarshal accepted invalid raw ledger") + } + return + } + if err != nil { + t.Fatalf("Unmarshal: %v", err) + } + keys, present := restored.WorkspaceEnrollmentBrokerKeys() + if present != tc.wantPresent || !reflect.DeepEqual(keys, tc.wantKeys) { + t.Fatalf("ledger = %q, %t; want %q, %t", keys, present, tc.wantKeys, tc.wantPresent) + } + }) + } +} diff --git a/engine/api/session.txt b/engine/api/session.txt index 7139d52830..add4aa269b 100644 --- a/engine/api/session.txt +++ b/engine/api/session.txt @@ -62,6 +62,7 @@ method func (*Session).RestoreSessionMetadata(kind SessionKind, rel SessionRelationship) error method func (*Session).RestoreTitleMetadata(title string, provenance TitleProvenance, revision uint64, generation TitleGenerationState, sources []string, attempts []TitleAttempt) method func (*Session).RestoreTokenUsage(usage map[UsageKind]TokenUsage) + method func (*Session).RestoreWorkspaceEnrollmentBrokerKeys(keys []string, present bool) error method func (*Session).ResumeWith() (PendingAsk, error) method func (*Session).RunID() string method func (*Session).SeedHistory(messages []Message) error @@ -77,6 +78,7 @@ method func (*Session).TokenUsageSnapshot() map[UsageKind]TokenUsage method func (*Session).UsageFor(kind UsageKind) Usage method func (*Session).ValidateAuthorizationState() error + method func (*Session).WorkspaceEnrollmentBrokerKeys() ([]string, bool) method func (Authority).Clone() Authority method func (Authority).Valid() bool method func (AuthorizationPayload).Valid() bool diff --git a/engine/session/session.go b/engine/session/session.go index 24ecb68355..0b7f7cd1c9 100644 --- a/engine/session/session.go +++ b/engine/session/session.go @@ -437,6 +437,12 @@ type Session struct { // A zero Authority with this marker is impossible: BindAuthority validates the // payload before setting either field. authorityBound bool + // workspaceEnrollmentBrokerKeys records exactly the broker registration keys + // installed by the last completed workspace enrollment. Presence distinguishes + // a fully restored empty ledger from unrepaired legacy state. + workspaceEnrollmentBrokerKeys []string + workspaceEnrollmentBrokerKeysPresent bool + workspaceEnrollmentBrokerKeysRestored bool // pending is set iff State == StateAwaiting. pending *PendingAsk // pendingAuthorization is set iff State == StateAuthorizing. diff --git a/engine/session/workspace_enrollment.go b/engine/session/workspace_enrollment.go index 252eb32fb7..d274b61017 100644 --- a/engine/session/workspace_enrollment.go +++ b/engine/session/workspace_enrollment.go @@ -2,6 +2,7 @@ package session import ( "fmt" + "sort" "time" "unicode" "unicode/utf8" @@ -117,9 +118,39 @@ func ValidWorkspaceEnrollmentToolNames(names []string) bool { return true } -// CompleteWorkspaceEnrollment atomically replaces the complete tool set after -// the exact pending enrollment has produced a verified catalogue. Authority is -// cloned from the aggregate so callers cannot supply or alter any non-tool axis. +// WorkspaceEnrollmentBrokerKeys returns an independent copy of the exact broker +// registration-key ledger and whether that provenance is present. Absent +// provenance is unrepaired legacy state, not a known empty ledger. +func (s *Session) WorkspaceEnrollmentBrokerKeys() ([]string, bool) { + return append([]string(nil), s.workspaceEnrollmentBrokerKeys...), s.workspaceEnrollmentBrokerKeysPresent +} + +// RestoreWorkspaceEnrollmentBrokerKeys restores copied workspace-enrollment +// provenance exactly once during root creation, snapshot restore, or successor +// construction. A present empty ledger is valid; absent provenance must have no +// keys. +func (s *Session) RestoreWorkspaceEnrollmentBrokerKeys(keys []string, present bool) error { + if s.workspaceEnrollmentBrokerKeysRestored { + return fmt.Errorf("%w: workspace enrollment broker keys already restored", ErrIllegalTransition) + } + if s.State != StateIdle || s.pending != nil || s.pendingAuthorization != nil || s.pendingWorkspaceEnrollment != nil { + return fmt.Errorf("%w: restore workspace enrollment broker keys requires an idle session", ErrIllegalTransition) + } + if (!present && len(keys) != 0) || !ValidWorkspaceEnrollmentToolNames(keys) { + return fmt.Errorf("%w: invalid workspace enrollment broker-key ledger", ErrIllegalTransition) + } + s.workspaceEnrollmentBrokerKeys = append([]string(nil), keys...) + s.workspaceEnrollmentBrokerKeysPresent = present + s.workspaceEnrollmentBrokerKeysRestored = true + return nil +} + +// CompleteWorkspaceEnrollment atomically replaces only the prior broker +// registration-key bundle after the exact pending enrollment has produced a +// verified catalogue. Authority is cloned from the aggregate so callers cannot +// supply or alter any non-tool axis. Previously recorded broker keys absent from +// carried authority remain excluded; keys newly introduced by this bundle remain +// eligible for authority. func (s *Session) CompleteWorkspaceEnrollment(pending PendingWorkspaceEnrollment, exactTools []string) error { if s.pendingWorkspaceEnrollment == nil || s.pendingWorkspaceEnrollment.ID != pending.ID || @@ -130,17 +161,79 @@ func (s *Session) CompleteWorkspaceEnrollment(pending PendingWorkspaceEnrollment if s.State != StateIdle || s.Conversation == nil { return fmt.Errorf("%w: workspace enrollment requires an idle session", ErrIllegalTransition) } - if !s.authorityBound || !ValidWorkspaceEnrollmentToolNames(exactTools) { + if !s.authorityBound || !s.workspaceEnrollmentBrokerKeysPresent || !ValidWorkspaceEnrollmentToolNames(exactTools) { return fmt.Errorf("%w: invalid workspace enrollment tool set", ErrIllegalTransition) } - exact := s.Authority.Clone() - exact.CapabilitySet.Tools = append([]string(nil), exactTools...) - s.Authority = exact + completedAuthority, ledger, err := completeWorkspaceEnrollmentAuthority( + s.Authority, s.workspaceEnrollmentBrokerKeys, exactTools, + ) + if err != nil { + return err + } + s.Authority = completedAuthority + s.workspaceEnrollmentBrokerKeys = ledger + s.workspaceEnrollmentBrokerKeysPresent = true s.pendingWorkspaceEnrollment = nil return nil } +func completeWorkspaceEnrollmentAuthority( + authority Authority, priorBrokerKeys, exactTools []string, +) (Authority, []string, error) { + priorBroker := make(map[string]struct{}, len(priorBrokerKeys)) + for _, name := range priorBrokerKeys { + priorBroker[name] = struct{}{} + } + priorAuthority := make(map[string]struct{}, len(authority.CapabilitySet.Tools)) + for _, name := range authority.CapabilitySet.Tools { + priorAuthority[name] = struct{}{} + } + excluded := make(map[string]struct{}, len(priorBroker)) + for name := range priorBroker { + if _, included := priorAuthority[name]; !included { + excluded[name] = struct{}{} + } + } + + completed := authority.Clone() + replacement := make([]string, 0, len(completed.CapabilitySet.Tools)+len(exactTools)) + replacementSet := make(map[string]struct{}, len(completed.CapabilitySet.Tools)+len(exactTools)) + for _, name := range completed.CapabilitySet.Tools { + if _, broker := priorBroker[name]; !broker { + replacement = append(replacement, name) + replacementSet[name] = struct{}{} + } + } + for _, name := range exactTools { + if _, excluded := excluded[name]; !excluded { + if _, alreadyPresent := replacementSet[name]; !alreadyPresent { + replacement = append(replacement, name) + replacementSet[name] = struct{}{} + } + } + } + ledger := append([]string(nil), exactTools...) + ledgerSet := make(map[string]struct{}, len(exactTools)) + for _, name := range exactTools { + ledgerSet[name] = struct{}{} + } + appendedExcluded := make([]string, 0, len(excluded)) + for name := range excluded { + if _, included := ledgerSet[name]; !included { + appendedExcluded = append(appendedExcluded, name) + } + } + sort.Strings(appendedExcluded) + ledger = append(ledger, appendedExcluded...) + if !ValidWorkspaceEnrollmentToolNames(ledger) { + return Authority{}, nil, fmt.Errorf("%w: invalid workspace enrollment broker-key ledger", ErrIllegalTransition) + } + + completed.CapabilitySet.Tools = replacement + return completed, ledger, nil +} + // AbortWorkspaceEnrollment clears only the enrollment whose exact identifier // is supplied. func (s *Session) AbortWorkspaceEnrollment(id WorkspaceEnrollmentID) error { diff --git a/engine/session/workspace_enrollment_authority_test.go b/engine/session/workspace_enrollment_authority_test.go new file mode 100644 index 0000000000..a31d17b06c --- /dev/null +++ b/engine/session/workspace_enrollment_authority_test.go @@ -0,0 +1,197 @@ +package session + +import ( + "errors" + "reflect" + "testing" +) + +func TestWorkspaceEnrollmentAuthority_Scenario1_ReplacesDynamicBrokerBundle(t *testing.T) { + s := newEnrollmentSession(t) + s.Authority.CapabilitySet.Tools = []string{"Read", "Memory", "broker-old", "broker-excluded"} + if err := s.RestoreWorkspaceEnrollmentBrokerKeys([]string{"broker-old", "broker-excluded"}, true); err != nil { + t.Fatalf("RestoreWorkspaceEnrollmentBrokerKeys: %v", err) + } + // Model an intentionally attenuated carried authority: the prior broker key is + // recorded but excluded, while unrelated capability absence remains untouched. + s.Authority.CapabilitySet.Tools = []string{"Read", "Memory", "broker-old"} + pending := testWorkspaceEnrollment() + if err := s.BeginWorkspaceEnrollment(pending); err != nil { + t.Fatalf("BeginWorkspaceEnrollment: %v", err) + } + if err := s.CompleteWorkspaceEnrollment(pending, []string{"broker-excluded", "broker-new"}); err != nil { + t.Fatalf("CompleteWorkspaceEnrollment: %v", err) + } + + got, bound := s.BoundAuthority() + if !bound { + t.Fatal("authority is not bound") + } + if want := []string{"Read", "Memory", "broker-new"}; !reflect.DeepEqual(got.CapabilitySet.Tools, want) { + t.Fatalf("authority tools = %q, want %q", got.CapabilitySet.Tools, want) + } + keys, present := s.WorkspaceEnrollmentBrokerKeys() + if !present || !reflect.DeepEqual(keys, []string{"broker-excluded", "broker-new"}) { + t.Fatalf("broker ledger = %q, %t; want completed bundle, true", keys, present) + } +} + +func TestWorkspaceEnrollmentAuthority_Scenario1_ExclusionSurvivesOmittingRefresh(t *testing.T) { + s := newEnrollmentSession(t) + s.Authority.CapabilitySet.Tools = []string{"Read", "A"} + if err := s.RestoreWorkspaceEnrollmentBrokerKeys([]string{"A", "C", "B"}, true); err != nil { + t.Fatalf("RestoreWorkspaceEnrollmentBrokerKeys: %v", err) + } + + pending := testWorkspaceEnrollment() + if err := s.BeginWorkspaceEnrollment(pending); err != nil { + t.Fatalf("BeginWorkspaceEnrollment: %v", err) + } + if err := s.CompleteWorkspaceEnrollment(pending, []string{"A"}); err != nil { + t.Fatalf("CompleteWorkspaceEnrollment omitting B: %v", err) + } + keys, present := s.WorkspaceEnrollmentBrokerKeys() + if !present || !reflect.DeepEqual(keys, []string{"A", "B", "C"}) { + t.Fatalf("ledger after omitting B and C = %q, %t; want [A B C], true", keys, present) + } + + if err := s.BeginWorkspaceEnrollment(pending); err != nil { + t.Fatalf("BeginWorkspaceEnrollment after refresh: %v", err) + } + if err := s.CompleteWorkspaceEnrollment(pending, []string{"A", "B"}); err != nil { + t.Fatalf("CompleteWorkspaceEnrollment reoffering B: %v", err) + } + got, _ := s.BoundAuthority() + if want := []string{"Read", "A"}; !reflect.DeepEqual(got.CapabilitySet.Tools, want) { + t.Fatalf("authority after reoffering B = %q, want %q", got.CapabilitySet.Tools, want) + } + + t.Run("overflow is atomic", func(t *testing.T) { + s := newEnrollmentSession(t) + excluded := make([]string, maxWorkspaceEnrollmentTools) + for i := range excluded { + excluded[i] = string(rune('a'+i/26)) + string(rune('a'+i%26)) + } + if err := s.RestoreWorkspaceEnrollmentBrokerKeys(excluded, true); err != nil { + t.Fatalf("RestoreWorkspaceEnrollmentBrokerKeys: %v", err) + } + pending := testWorkspaceEnrollment() + if err := s.BeginWorkspaceEnrollment(pending); err != nil { + t.Fatalf("BeginWorkspaceEnrollment: %v", err) + } + beforeAuthority, _ := s.BoundAuthority() + beforeKeys, beforePresent := s.WorkspaceEnrollmentBrokerKeys() + if err := s.CompleteWorkspaceEnrollment(pending, []string{"new"}); !errors.Is(err, ErrIllegalTransition) { + t.Fatalf("CompleteWorkspaceEnrollment overflow error = %v, want ErrIllegalTransition", err) + } + afterAuthority, _ := s.BoundAuthority() + afterKeys, afterPresent := s.WorkspaceEnrollmentBrokerKeys() + if !reflect.DeepEqual(afterAuthority, beforeAuthority) || !reflect.DeepEqual(afterKeys, beforeKeys) || afterPresent != beforePresent { + t.Fatalf("overflow mutated authority or ledger: authority=%#v keys=%q present=%t", afterAuthority, afterKeys, afterPresent) + } + if got, ok := s.PendingWorkspaceEnrollment(); !ok || got != pending { + t.Fatalf("overflow changed pending enrollment: %+v, %t", got, ok) + } + }) +} + +func TestWorkspaceEnrollmentAuthority_Scenario1_PreservesUnrelatedCapabilities(t *testing.T) { + s := newEnrollmentSession(t) + if err := s.RestoreWorkspaceEnrollmentBrokerKeys(nil, true); err != nil { + t.Fatalf("RestoreWorkspaceEnrollmentBrokerKeys: %v", err) + } + s.Authority.CapabilitySet.Tools = []string{"Read", "Memory"} + pending := testWorkspaceEnrollment() + if err := s.BeginWorkspaceEnrollment(pending); err != nil { + t.Fatal(err) + } + if err := s.CompleteWorkspaceEnrollment(pending, []string{"broker-new"}); err != nil { + t.Fatal(err) + } + got, _ := s.BoundAuthority() + if want := []string{"Read", "Memory", "broker-new"}; !reflect.DeepEqual(got.CapabilitySet.Tools, want) { + t.Fatalf("authority tools = %q, want %q", got.CapabilitySet.Tools, want) + } +} + +func TestWorkspaceEnrollmentAuthority_Scenario1_DoesNotWidenAttenuatedAuthority(t *testing.T) { + s := newEnrollmentSession(t) + s.Authority.CapabilitySet.Tools = []string{"Read"} + if err := s.RestoreWorkspaceEnrollmentBrokerKeys([]string{"broker-old"}, true); err != nil { + t.Fatal(err) + } + pending := testWorkspaceEnrollment() + if err := s.BeginWorkspaceEnrollment(pending); err != nil { + t.Fatal(err) + } + if err := s.CompleteWorkspaceEnrollment(pending, []string{"broker-old", "broker-new"}); err != nil { + t.Fatal(err) + } + got, _ := s.BoundAuthority() + if want := []string{"Read", "broker-new"}; !reflect.DeepEqual(got.CapabilitySet.Tools, want) { + t.Fatalf("authority tools = %q, want %q", got.CapabilitySet.Tools, want) + } +} + +func TestWorkspaceEnrollmentAuthority_Scenario1_PreservesSupportedRestrictions(t *testing.T) { + s := newEnrollmentSession(t) + if err := s.RestoreWorkspaceEnrollmentBrokerKeys(nil, true); err != nil { + t.Fatalf("RestoreWorkspaceEnrollmentBrokerKeys: %v", err) + } + before := s.Authority.Clone() + s.Authority.CapabilitySet.Tools = []string{"Read"} + before.CapabilitySet.Tools = []string{"Read"} + pending := testWorkspaceEnrollment() + if err := s.BeginWorkspaceEnrollment(pending); err != nil { + t.Fatal(err) + } + if err := s.CompleteWorkspaceEnrollment(pending, []string{"broker"}); err != nil { + t.Fatal(err) + } + got, _ := s.BoundAuthority() + got.CapabilitySet.Tools = nil + before.CapabilitySet.Tools = nil + if !reflect.DeepEqual(got, before) { + t.Fatalf("completion changed non-tool authority axes: got %#v, want %#v", got, before) + } +} + +func TestWorkspaceEnrollmentAuthority_Scenario2_RefreshSubtractsRestoredBrokerKeys(t *testing.T) { + s := newEnrollmentSession(t) + s.Authority.CapabilitySet.Tools = []string{"Read", "broker-old"} + if err := s.RestoreWorkspaceEnrollmentBrokerKeys([]string{"broker-old"}, true); err != nil { + t.Fatal(err) + } + pending := testWorkspaceEnrollment() + if err := s.BeginWorkspaceEnrollment(pending); err != nil { + t.Fatal(err) + } + if err := s.CompleteWorkspaceEnrollment(pending, []string{"broker-new"}); err != nil { + t.Fatal(err) + } + got, _ := s.BoundAuthority() + if want := []string{"Read", "broker-new"}; !reflect.DeepEqual(got.CapabilitySet.Tools, want) { + t.Fatalf("authority after restored-ledger refresh = %q, want %q", got.CapabilitySet.Tools, want) + } +} +func TestWorkspaceEnrollmentAuthority_Scenario2_ProvenancePresenceAndBounds(t *testing.T) { + s := newEnrollmentSession(t) + keys, present := s.WorkspaceEnrollmentBrokerKeys() + if present || len(keys) != 0 { + t.Fatalf("new session ledger = %q, %t; want absent", keys, present) + } + if err := s.RestoreWorkspaceEnrollmentBrokerKeys(nil, true); err != nil { + t.Fatalf("RestoreWorkspaceEnrollmentBrokerKeys: %v", err) + } + if err := s.RestoreWorkspaceEnrollmentBrokerKeys(nil, true); !errors.Is(err, ErrIllegalTransition) { + t.Fatalf("second RestoreWorkspaceEnrollmentBrokerKeys error = %v, want ErrIllegalTransition", err) + } + + s = newEnrollmentSession(t) + if err := s.RestoreWorkspaceEnrollmentBrokerKeys([]string{"broker"}, false); err == nil { + t.Fatal("accepted keys with absent provenance") + } + if err := s.RestoreWorkspaceEnrollmentBrokerKeys([]string{"broker", "broker"}, true); err == nil { + t.Fatal("accepted duplicate broker keys") + } +} diff --git a/engine/session/workspace_enrollment_test.go b/engine/session/workspace_enrollment_test.go index f104d77d17..ac1a6318e0 100644 --- a/engine/session/workspace_enrollment_test.go +++ b/engine/session/workspace_enrollment_test.go @@ -248,6 +248,9 @@ func sessionWithEnrollmentAuthority(t *testing.T) (*Session, PendingWorkspaceEnr if err := s.BindAuthority(enrollmentAuthority("Read", "stale")); err != nil { t.Fatalf("BindAuthority: %v", err) } + if err := s.RestoreWorkspaceEnrollmentBrokerKeys([]string{"Read", "stale"}, true); err != nil { + t.Fatalf("RestoreWorkspaceEnrollmentBrokerKeys: %v", err) + } pending := testWorkspaceEnrollment() if err := s.BeginWorkspaceEnrollment(pending); err != nil { t.Fatalf("BeginWorkspaceEnrollment: %v", err) diff --git a/internal/adapter/mcpbroker/query_continuation_test.go b/internal/adapter/mcpbroker/query_continuation_test.go index ea51ca7290..7c48d1a551 100644 --- a/internal/adapter/mcpbroker/query_continuation_test.go +++ b/internal/adapter/mcpbroker/query_continuation_test.go @@ -85,7 +85,7 @@ func TestCallMcpWithQueryBrokerSupport_Scenario1_AuthorizationContinuation(t *te defer runtime.Close() store := memstore.New() buildHost := func(first bool) *server.Service { - factory := func(_ context.Context, _ server.ProviderSelector, _ []mcp.ServerConfig, _ server.SessionProfile, _ string, mode session.PermissionMode, tools []tool.Tool) (server.SessionEngineResult, error) { + factory := func(_ context.Context, _ server.ProviderSelector, _ []mcp.ServerConfig, _ server.SessionProfile, _ string, mode session.PermissionMode, tools []tool.Tool, toolKeys []string) (server.SessionEngineResult, error) { cat := tool.NewCatalog() for _, candidate := range tools { cat.MustRegister(candidate) @@ -94,10 +94,10 @@ func TestCallMcpWithQueryBrokerSupport_Scenario1_AuthorizationContinuation(t *te if first { turns = append([]mockllm.Turn{mockllm.ToolCallTurn(call)}, turns...) } - return server.SessionEngineResult{Engine: agent.NewEngine(agent.Deps{LLM: mockllm.New(turns...), Catalog: cat, Policy: queryAllowPolicy{}, Store: store}), BuiltForMode: mode, Close: func() error { return nil }}, nil + return server.SessionEngineResult{Engine: agent.NewEngine(agent.Deps{LLM: mockllm.New(turns...), Catalog: cat, Policy: queryAllowPolicy{}, Store: store}), BuiltForMode: mode, BrokerRegistrationKeys: append([]string(nil), toolKeys...), Close: func() error { return nil }}, nil } svc, err := server.NewService(server.Config{Engine: agent.NewEngine(agent.Deps{LLM: mockllm.New(mockllm.TextTurn("done")), Catalog: tool.NewCatalog(), Policy: queryAllowPolicy{}}), Store: store, MCPBroker: runtime, SessionEngineWithTools: factory, SessionEngine: func(ctx context.Context, sel server.ProviderSelector, specs []mcp.ServerConfig, profile server.SessionProfile, workspace string, mode session.PermissionMode) (server.SessionEngineResult, error) { - return factory(ctx, sel, specs, profile, workspace, mode, nil) + return factory(ctx, sel, specs, profile, workspace, mode, nil, nil) }, PlacementProvider: queryPlacement{}, PlacementScope: "test", NewID: func() session.SessionID { return "authorization-session" }}) if err != nil { t.Fatal(err) diff --git a/internal/adapter/server/authorization_close_test.go b/internal/adapter/server/authorization_close_test.go index ece29822f6..f6a0e1ff51 100644 --- a/internal/adapter/server/authorization_close_test.go +++ b/internal/adapter/server/authorization_close_test.go @@ -100,7 +100,7 @@ func TestCloseCompletesWhenOneSessionSettlementFails(t *testing.T) { RootAuthority: func(session.SessionKind) session.Authority { return session.Authority{CapabilitySet: governance.CapabilitySet{}, Provenance: "test"} }, - SessionEngineWithTools: func(context.Context, ProviderSelector, []mcp.ServerConfig, SessionProfile, string, session.PermissionMode, []tool.Tool) (SessionEngineResult, error) { + SessionEngineWithTools: func(context.Context, ProviderSelector, []mcp.ServerConfig, SessionProfile, string, session.PermissionMode, []tool.Tool, []string) (SessionEngineResult, error) { return brokerEngineResult(), nil }, }) diff --git a/internal/adapter/server/broker_rebind_test.go b/internal/adapter/server/broker_rebind_test.go index b036345a32..aa3ce961ec 100644 --- a/internal/adapter/server/broker_rebind_test.go +++ b/internal/adapter/server/broker_rebind_test.go @@ -32,7 +32,7 @@ func TestWorkspaceEnrollmentRebindsAfterBrokerRestart(t *testing.T) { RootAuthority: func(session.SessionKind) session.Authority { return session.Authority{CapabilitySet: governance.CapabilitySet{Tools: []string{"mcp__calendar__list"}}, Provenance: "test"} }, - SessionEngineWithTools: func(context.Context, ProviderSelector, []mcp.ServerConfig, SessionProfile, string, session.PermissionMode, []tool.Tool) (SessionEngineResult, error) { + SessionEngineWithTools: func(context.Context, ProviderSelector, []mcp.ServerConfig, SessionProfile, string, session.PermissionMode, []tool.Tool, []string) (SessionEngineResult, error) { return brokerEngineResult(), nil }, }) diff --git a/internal/adapter/server/harness_context_test.go b/internal/adapter/server/harness_context_test.go index 09c8de2736..807c2704ec 100644 --- a/internal/adapter/server/harness_context_test.go +++ b/internal/adapter/server/harness_context_test.go @@ -44,7 +44,7 @@ func TestHarnessContextGeneratedIDCollisionAfterClose(t *testing.T) { t.Fatal("legacy factory called") return SessionEngineResult{}, nil }, - SessionContextEngine: func(context.Context, session.SessionID, *session.Principal, ExecutionWorkspaceAcquirer, ProviderSelector, []mcp.ServerConfig, SessionProfile, string, session.PermissionMode, []tool.Tool, session.SessionID) (SessionEngineResult, error) { + SessionContextEngine: func(context.Context, session.SessionID, *session.Principal, ExecutionWorkspaceAcquirer, ProviderSelector, []mcp.ServerConfig, SessionProfile, string, session.PermissionMode, []tool.Tool, []string, session.SessionID) (SessionEngineResult, error) { builds++ return SessionEngineResult{Engine: eng, Close: func() error { closes++; return nil }}, nil }, @@ -127,7 +127,7 @@ func TestHarnessContextGeneratedIDPublicationCollision(t *testing.T) { t.Fatal("legacy factory called") return SessionEngineResult{}, nil }, - SessionContextEngine: func(ctx context.Context, id session.SessionID, _ *session.Principal, _ ExecutionWorkspaceAcquirer, _ ProviderSelector, _ []mcp.ServerConfig, _ SessionProfile, _ string, _ session.PermissionMode, _ []tool.Tool, _ session.SessionID) (SessionEngineResult, error) { + SessionContextEngine: func(ctx context.Context, id session.SessionID, _ *session.Principal, _ ExecutionWorkspaceAcquirer, _ ProviderSelector, _ []mcp.ServerConfig, _ SessionProfile, _ string, _ session.PermissionMode, _ []tool.Tool, _ []string, _ session.SessionID) (SessionEngineResult, error) { // The losing Service has already observed absence and reserved its ID. var err error winner, err = winnerService.CreateSessionWithProfile(ctx, tc.mode, session.Limits{}, ProviderSelector{}, ProfileDefault, WithSessionID(id), WithOwner(tc.winnerOwner)) @@ -189,7 +189,7 @@ func TestHarnessContextUnpublishedCreateRetiresBinding(t *testing.T) { t.Fatal("legacy factory called") return SessionEngineResult{}, nil }, - SessionContextEngine: func(_ context.Context, id session.SessionID, _ *session.Principal, _ ExecutionWorkspaceAcquirer, _ ProviderSelector, _ []mcp.ServerConfig, _ SessionProfile, _ string, _ session.PermissionMode, _ []tool.Tool, _ session.SessionID) (SessionEngineResult, error) { + SessionContextEngine: func(_ context.Context, id session.SessionID, _ *session.Principal, _ ExecutionWorkspaceAcquirer, _ ProviderSelector, _ []mcp.ServerConfig, _ SessionProfile, _ string, _ session.PermissionMode, _ []tool.Tool, _ []string, _ session.SessionID) (SessionEngineResult, error) { bound = id return SessionEngineResult{Engine: eng, Close: func() error { closed = true; return nil }}, nil }, diff --git a/internal/adapter/server/mcp_authorization.go b/internal/adapter/server/mcp_authorization.go index cf2a344140..73ae7c0969 100644 --- a/internal/adapter/server/mcp_authorization.go +++ b/internal/adapter/server/mcp_authorization.go @@ -417,7 +417,7 @@ func (s *Service) continueGrantedAuthorizationLocked(ctx context.Context, sess * // authenticated snapshot and compensates a failed handoff by restoring its claim. func (s *Service) rebuildGrantedAuthorizationEngine(ctx context.Context, sess *session.Session, claimed session.PendingAuthorization, exactTools []tool.Tool) error { sel := ProviderSelector{ProviderID: sess.ProviderID, ModelID: sess.ModelID, ReasoningEffort: sess.ReasoningEffort} - if _, err := s.buildAndRegisterSessionEngineWithBrokerTools(ctx, sess, sel, profileForSession(sess), sess.Mode, true, exactTools, true); err != nil { + if _, err := s.buildAndRegisterSessionEngineWithBrokerTools(ctx, sess, sel, profileForSession(sess), sess.Mode, true, exactTools, nil, true); err != nil { if errors.Is(err, ErrInvalidArgument) { // Some OTHER run raced onto this session's slot between the parked run's // admission and this rebuild — the same race registerAndStartGrantedAuthorization diff --git a/internal/adapter/server/mcp_authorization_lifecycle_test.go b/internal/adapter/server/mcp_authorization_lifecycle_test.go index e120920abd..cc5ca6bdbf 100644 --- a/internal/adapter/server/mcp_authorization_lifecycle_test.go +++ b/internal/adapter/server/mcp_authorization_lifecycle_test.go @@ -344,7 +344,7 @@ func newLifecycleFixtureConfigured(t *testing.T, status session.AuthorizationSta cfg.SessionEngine = func(_ context.Context, _ ProviderSelector, _ []mcp.ServerConfig, _ SessionProfile, _ string, mode session.PermissionMode) (SessionEngineResult, error) { return SessionEngineResult{Engine: buildEngine(nil), BuiltForMode: mode, Close: func() error { return nil }}, nil } - cfg.SessionEngineWithTools = func(_ context.Context, _ ProviderSelector, specs []mcp.ServerConfig, _ SessionProfile, _ string, mode session.PermissionMode, tools []tool.Tool) (SessionEngineResult, error) { + cfg.SessionEngineWithTools = func(_ context.Context, _ ProviderSelector, specs []mcp.ServerConfig, _ SessionProfile, _ string, mode session.PermissionMode, tools []tool.Tool, _ []string) (SessionEngineResult, error) { names := make([]string, 0, len(tools)) for _, candidate := range tools { names = append(names, candidate.Spec().Name) @@ -1580,7 +1580,7 @@ func TestMCPAuthorizationUnpinnedEngineResolutionKeepsStartupRevision(t *testing return SessionEngineResult{Engine: f.svc.cfg.Engine, BuiltForMode: mode, RuntimeRevision: revision, Close: func() error { return nil }}, nil } f.svc.cfg.SessionEngine = build - f.svc.cfg.SessionEngineWithTools = func(ctx context.Context, sel ProviderSelector, specs []mcp.ServerConfig, profile SessionProfile, workspace string, mode session.PermissionMode, _ []tool.Tool) (SessionEngineResult, error) { + f.svc.cfg.SessionEngineWithTools = func(ctx context.Context, sel ProviderSelector, specs []mcp.ServerConfig, profile SessionProfile, workspace string, mode session.PermissionMode, _ []tool.Tool, _ []string) (SessionEngineResult, error) { return build(ctx, sel, specs, profile, workspace, mode) } @@ -1622,7 +1622,7 @@ func TestMCPAuthorizationUnpinnedEngineResolutionKeepsStartupRevision(t *testing return SessionEngineResult{Engine: withoutRuntime.svc.cfg.Engine, BuiltForMode: mode, Close: func() error { return nil }}, nil } withoutRuntime.svc.cfg.SessionEngine = withoutBuild - withoutRuntime.svc.cfg.SessionEngineWithTools = func(ctx context.Context, sel ProviderSelector, specs []mcp.ServerConfig, profile SessionProfile, workspace string, mode session.PermissionMode, _ []tool.Tool) (SessionEngineResult, error) { + withoutRuntime.svc.cfg.SessionEngineWithTools = func(ctx context.Context, sel ProviderSelector, specs []mcp.ServerConfig, profile SessionProfile, workspace string, mode session.PermissionMode, _ []tool.Tool, _ []string) (SessionEngineResult, error) { return withoutBuild(ctx, sel, specs, profile, workspace, mode) } for range 2 { diff --git a/internal/adapter/server/mcp_authorization_transport_test.go b/internal/adapter/server/mcp_authorization_transport_test.go index b9bad74613..67538c15d6 100644 --- a/internal/adapter/server/mcp_authorization_transport_test.go +++ b/internal/adapter/server/mcp_authorization_transport_test.go @@ -241,7 +241,7 @@ func TestMCPAuthorizationGRPCStaleApprovalRefusesThenAcceptsCorrectRetry(t *test followup := session.NewToolCall("followup-stale", "protected", nil) f := newLifecycleFixtureWithTurns(t, session.AuthorizationGranted, nil, time.Now, nil, mockllm.ToolCallTurn(followup), mockllm.TextTurn("continued after stale retry")) - f.svc.cfg.SessionEngineWithTools = func(_ context.Context, _ ProviderSelector, _ []mcp.ServerConfig, _ SessionProfile, _ string, mode session.PermissionMode, tools []tool.Tool) (SessionEngineResult, error) { + f.svc.cfg.SessionEngineWithTools = func(_ context.Context, _ ProviderSelector, _ []mcp.ServerConfig, _ SessionProfile, _ string, mode session.PermissionMode, tools []tool.Tool, _ []string) (SessionEngineResult, error) { catalog := tool.NewCatalog() for _, candidate := range tools { catalog.MustRegister(candidate) diff --git a/internal/adapter/server/mcp_broker.go b/internal/adapter/server/mcp_broker.go index bd5a7c5beb..ef18de3206 100644 --- a/internal/adapter/server/mcp_broker.go +++ b/internal/adapter/server/mcp_broker.go @@ -4,6 +4,7 @@ import ( "context" "errors" "fmt" + "slices" "github.com/stacklok/mecatl/engine/port" "github.com/stacklok/mecatl/engine/session" @@ -47,19 +48,19 @@ func brokerTools(local *localBrokerAttachment) []tool.Tool { return tools } -func (s *Service) callSessionEngine(ctx context.Context, id session.SessionID, owner *session.Principal, acquire ExecutionWorkspaceAcquirer, sel ProviderSelector, specs []mcp.ServerConfig, profile SessionProfile, workspace string, mode session.PermissionMode, sessionTools []tool.Tool, forkSource ...session.SessionID) (SessionEngineResult, error) { +func (s *Service) callSessionEngine(ctx context.Context, id session.SessionID, owner *session.Principal, acquire ExecutionWorkspaceAcquirer, sel ProviderSelector, specs []mcp.ServerConfig, profile SessionProfile, workspace string, mode session.PermissionMode, sessionTools []tool.Tool, sessionToolKeys []string, forkSource ...session.SessionID) (SessionEngineResult, error) { if s.cfg.SessionContextEngine != nil { var source session.SessionID if len(forkSource) > 0 { source = forkSource[0] } - return s.cfg.SessionContextEngine(ctx, id, owner.Clone(), acquire, sel, specs, profile, workspace, mode, append([]tool.Tool(nil), sessionTools...), source) + return s.cfg.SessionContextEngine(ctx, id, owner.Clone(), acquire, sel, specs, profile, workspace, mode, append([]tool.Tool(nil), sessionTools...), slices.Clone(sessionToolKeys), source) } if s.cfg.MCPBroker != nil { if s.cfg.SessionEngineWithTools == nil { return SessionEngineResult{}, fmt.Errorf("%w: broker tools require an explicit session catalogue factory", ErrConfig) } - return s.cfg.SessionEngineWithTools(ctx, sel, specs, profile, workspace, mode, append([]tool.Tool(nil), sessionTools...)) + return s.cfg.SessionEngineWithTools(ctx, sel, specs, profile, workspace, mode, append([]tool.Tool(nil), sessionTools...), slices.Clone(sessionToolKeys)) } if s.cfg.SessionEngine == nil { return SessionEngineResult{}, fmt.Errorf("%w: per-session engine not supported (no session-engine factory configured)", ErrInvalidArgument) diff --git a/internal/adapter/server/mcp_broker_id_validation_test.go b/internal/adapter/server/mcp_broker_id_validation_test.go index 01366bcf9a..f36e1f04d8 100644 --- a/internal/adapter/server/mcp_broker_id_validation_test.go +++ b/internal/adapter/server/mcp_broker_id_validation_test.go @@ -52,7 +52,7 @@ func TestCreateSessionWithBrokerInvalidIDRejected(t *testing.T) { factoryCalls++ return brokerEngineResult(), nil }, - SessionEngineWithTools: func(context.Context, ProviderSelector, []mcp.ServerConfig, SessionProfile, string, session.PermissionMode, []tool.Tool) (SessionEngineResult, error) { + SessionEngineWithTools: func(context.Context, ProviderSelector, []mcp.ServerConfig, SessionProfile, string, session.PermissionMode, []tool.Tool, []string) (SessionEngineResult, error) { factoryCalls++ return brokerEngineResult(), nil }, diff --git a/internal/adapter/server/mcp_broker_multi_upstream_e2e_test.go b/internal/adapter/server/mcp_broker_multi_upstream_e2e_test.go index c7fb619290..cb05b24eb0 100644 --- a/internal/adapter/server/mcp_broker_multi_upstream_e2e_test.go +++ b/internal/adapter/server/mcp_broker_multi_upstream_e2e_test.go @@ -13,6 +13,7 @@ import ( "net/http/httptest" "net/url" "reflect" + "sort" "strings" "sync" "testing" @@ -262,7 +263,7 @@ func newMultiUpstreamFixture(t *testing.T, rejectB bool, contexts ...context.Con RootAuthority: func(session.SessionKind) session.Authority { return session.Authority{CapabilitySet: governance.CapabilitySet{Tools: []string{"mcp__backend-a__whoami", "mcp__backend-b__whoami"}}, Provenance: "broker-e2e-test"} }, - SessionEngineWithTools: func(_ context.Context, _ ProviderSelector, _ []mcp.ServerConfig, _ SessionProfile, _ string, _ session.PermissionMode, tools []tool.Tool) (SessionEngineResult, error) { + SessionEngineWithTools: func(_ context.Context, _ ProviderSelector, _ []mcp.ServerConfig, _ SessionProfile, _ string, _ session.PermissionMode, tools []tool.Tool, keys []string) (SessionEngineResult, error) { f.mu.Lock() defer f.mu.Unlock() f.catalogueBuilds = append(f.catalogueBuilds, append([]tool.Tool(nil), tools...)) @@ -270,7 +271,11 @@ func newMultiUpstreamFixture(t *testing.T, rejectB bool, contexts ...context.Con f.failBuilds-- return SessionEngineResult{}, errors.New("injected engine replacement failure") } - return brokerEngineResult(), nil + return SessionEngineResult{ + Engine: brokerEngineResult().Engine, + BrokerRegistrationKeys: append([]string(nil), keys...), + NonBrokerRegistrationKeys: []string{"Read"}, + }, nil }, }) if err != nil { @@ -292,6 +297,15 @@ func newMultiUpstreamFixture(t *testing.T, rejectB bool, contexts ...context.Con return f } +func toolNames(tools []tool.Tool) []string { + names := make([]string, 0, len(tools)) + for _, registered := range tools { + names = append(names, registered.Spec().Name) + } + sort.Strings(names) + return names +} + func (f *multiUpstreamFixture) start() WorkspaceEnrollmentProjection { f.t.Helper() started, err := f.service.ConnectWorkspaceServices(f.t.Context(), f.session.ID) diff --git a/internal/adapter/server/mcp_broker_test.go b/internal/adapter/server/mcp_broker_test.go index a54e99e69c..df030eae10 100644 --- a/internal/adapter/server/mcp_broker_test.go +++ b/internal/adapter/server/mcp_broker_test.go @@ -47,18 +47,35 @@ func brokerEngineResult() SessionEngineResult { return SessionEngineResult{Engine: agent.NewEngine(agent.Deps{LLM: mockllm.New(mockllm.TextTurn("done")), Catalog: tool.NewCatalog(), Policy: permpolicy.NewPolicy(nil, nil)}), Close: func() error { return nil }} } +func brokerEngineResultWithTools(tools []tool.Tool) SessionEngineResult { + result := brokerEngineResult() + result.BrokerRegistrationKeys = toolNames(tools) + result.NonBrokerRegistrationKeys = []string{"Read"} + return result +} + // brokerEngineResultWithStore mirrors brokerEngineResult but wires the // engine's own Store: brokerEngineResult's engine has none, which is a silent // no-op save for tests that only exercise enrollment controls, but a test // that drives a REAL StartRunContent turn needs it wired to the session's // actual store, or nothing the run does ever persists. func brokerEngineResultWithStore(store port.SessionStore) SessionEngineResult { + return brokerEngineResultWithStoreAndTools(store, nil) +} + +// brokerEngineResultWithStoreAndTools reports the exact registration metadata +// for the tools supplied to the test factory. The fixture catalog's remaining +// non-broker registration is Read; keeping that partition explicit exercises +// the production handoff instead of relying on an omitted-field fallback. +func brokerEngineResultWithStoreAndTools(store port.SessionStore, tools []tool.Tool) SessionEngineResult { return SessionEngineResult{ Engine: agent.NewEngine(agent.Deps{ LLM: mockllm.New(mockllm.TextTurn("done")), Catalog: tool.NewCatalog(), Policy: permpolicy.NewPolicy(nil, nil), Store: store, }), - Close: func() error { return nil }, + BrokerRegistrationKeys: toolNames(tools), + NonBrokerRegistrationKeys: []string{"Read"}, + Close: func() error { return nil }, } } @@ -97,7 +114,7 @@ func TestMCPBrokerCarryoverDoesNotWidenAttenuatedAuthority(t *testing.T) { RootAuthority: func(session.SessionKind) session.Authority { return session.Authority{CapabilitySet: governance.CapabilitySet{Tools: []string{"root-only"}}, Provenance: "fresh-root"} }, - SessionEngineWithTools: func(context.Context, ProviderSelector, []mcp.ServerConfig, SessionProfile, string, session.PermissionMode, []tool.Tool) (SessionEngineResult, error) { + SessionEngineWithTools: func(context.Context, ProviderSelector, []mcp.ServerConfig, SessionProfile, string, session.PermissionMode, []tool.Tool, []string) (SessionEngineResult, error) { return brokerEngineResult(), nil }, }) @@ -129,7 +146,7 @@ func TestMCPBrokerCanonicalAttachPersistReattachAndLocalClose(t *testing.T) { t.Fatal("broker session used the factory without explicit tools") return SessionEngineResult{}, nil }, - SessionEngineWithTools: func(_ context.Context, _ ProviderSelector, _ []mcp.ServerConfig, _ SessionProfile, _ string, _ session.PermissionMode, tools []tool.Tool) (SessionEngineResult, error) { + SessionEngineWithTools: func(_ context.Context, _ ProviderSelector, _ []mcp.ServerConfig, _ SessionProfile, _ string, _ session.PermissionMode, tools []tool.Tool, _ []string) (SessionEngineResult, error) { factoryCalls++ if tools == nil { t.Fatal("explicit broker tool slice was nil") @@ -191,7 +208,7 @@ func TestMCPBrokerPermanentDeleteRemovesLogicalState(t *testing.T) { Engine: brokerEngineResult().Engine, Store: memstore.New(), PlacementProvider: brokerPlacementProvider{}, PlacementScope: "test", NewID: func() session.SessionID { return "permanent-delete" }, MCPBroker: runtime, - SessionEngineWithTools: func(context.Context, ProviderSelector, []mcp.ServerConfig, SessionProfile, string, session.PermissionMode, []tool.Tool) (SessionEngineResult, error) { + SessionEngineWithTools: func(context.Context, ProviderSelector, []mcp.ServerConfig, SessionProfile, string, session.PermissionMode, []tool.Tool, []string) (SessionEngineResult, error) { return brokerEngineResult(), nil }, }) @@ -262,7 +279,7 @@ func newFlakyDeleteService(t *testing.T, id session.SessionID) (*Service, *adapt Engine: brokerEngineResult().Engine, Store: store, PlacementProvider: brokerPlacementProvider{}, PlacementScope: "test", NewID: func() session.SessionID { return id }, MCPBroker: broker, - SessionEngineWithTools: func(context.Context, ProviderSelector, []mcp.ServerConfig, SessionProfile, string, session.PermissionMode, []tool.Tool) (SessionEngineResult, error) { + SessionEngineWithTools: func(context.Context, ProviderSelector, []mcp.ServerConfig, SessionProfile, string, session.PermissionMode, []tool.Tool, []string) (SessionEngineResult, error) { return brokerEngineResult(), nil }, }) @@ -348,7 +365,7 @@ func TestDeleteSessionKeepsBrokerStateWhenDurableDeleteFails(t *testing.T) { Engine: brokerEngineResult().Engine, Store: store, PlacementProvider: brokerPlacementProvider{}, PlacementScope: "test", NewID: func() session.SessionID { return "keep-broker-state" }, MCPBroker: runtime, - SessionEngineWithTools: func(context.Context, ProviderSelector, []mcp.ServerConfig, SessionProfile, string, session.PermissionMode, []tool.Tool) (SessionEngineResult, error) { + SessionEngineWithTools: func(context.Context, ProviderSelector, []mcp.ServerConfig, SessionProfile, string, session.PermissionMode, []tool.Tool, []string) (SessionEngineResult, error) { return brokerEngineResult(), nil }, }) @@ -395,7 +412,7 @@ func TestMCPBrokerAttachRollbackSerializesAgainstReattachCommit(t *testing.T) { Engine: brokerEngineResult().Engine, Store: store, PlacementProvider: brokerPlacementProvider{}, PlacementScope: "test", NewID: func() session.SessionID { return "rollback-race" }, MCPBroker: runtime, - SessionEngineWithTools: func(context.Context, ProviderSelector, []mcp.ServerConfig, SessionProfile, string, session.PermissionMode, []tool.Tool) (SessionEngineResult, error) { + SessionEngineWithTools: func(context.Context, ProviderSelector, []mcp.ServerConfig, SessionProfile, string, session.PermissionMode, []tool.Tool, []string) (SessionEngineResult, error) { return brokerEngineResult(), nil }, }) @@ -455,7 +472,7 @@ func TestMCPBrokerCloseSerializesAgainstEngineRebuild(t *testing.T) { Engine: brokerEngineResult().Engine, Store: memstore.New(), PlacementProvider: brokerPlacementProvider{}, PlacementScope: "test", NewID: func() session.SessionID { return "close-race" }, MCPBroker: runtime, - SessionEngineWithTools: func(_ context.Context, _ ProviderSelector, _ []mcp.ServerConfig, _ SessionProfile, _ string, _ session.PermissionMode, tools []tool.Tool) (SessionEngineResult, error) { + SessionEngineWithTools: func(_ context.Context, _ ProviderSelector, _ []mcp.ServerConfig, _ SessionProfile, _ string, _ session.PermissionMode, tools []tool.Tool, _ []string) (SessionEngineResult, error) { factoryCalls.Add(1) wrapper = tools[0] return brokerEngineResult(), nil @@ -544,7 +561,7 @@ func TestMCPBrokerServiceShutdownBoundsAttachmentDrain(t *testing.T) { Engine: brokerEngineResult().Engine, Store: memstore.New(), Diagnostics: port.NopDiagnostics{}, PlacementProvider: brokerPlacementProvider{}, PlacementScope: "test", NewID: func() session.SessionID { return "shutdown-drain" }, MCPBroker: runtime, - SessionEngineWithTools: func(_ context.Context, _ ProviderSelector, _ []mcp.ServerConfig, _ SessionProfile, _ string, _ session.PermissionMode, tools []tool.Tool) (SessionEngineResult, error) { + SessionEngineWithTools: func(_ context.Context, _ ProviderSelector, _ []mcp.ServerConfig, _ SessionProfile, _ string, _ session.PermissionMode, tools []tool.Tool, _ []string) (SessionEngineResult, error) { wrapper = tools[0] return brokerEngineResult(), nil }, @@ -585,7 +602,7 @@ func TestMCPBrokerCreateRollbackDeletesOnlyNewLogicalState(t *testing.T) { PlacementProvider: brokerPlacementProvider{}, PlacementScope: "test", NewID: func() session.SessionID { return "rollback" }, MCPBroker: runtime, - SessionEngineWithTools: func(context.Context, ProviderSelector, []mcp.ServerConfig, SessionProfile, string, session.PermissionMode, []tool.Tool) (SessionEngineResult, error) { + SessionEngineWithTools: func(context.Context, ProviderSelector, []mcp.ServerConfig, SessionProfile, string, session.PermissionMode, []tool.Tool, []string) (SessionEngineResult, error) { return brokerEngineResult(), nil }, }) diff --git a/internal/adapter/server/placement_successor.go b/internal/adapter/server/placement_successor.go index 788c6ea113..0db6379786 100644 --- a/internal/adapter/server/placement_successor.go +++ b/internal/adapter/server/placement_successor.go @@ -202,11 +202,12 @@ func (s *Service) createPlacedSuccessorLocked(ctx context.Context, req ForkSucce created := session.New(destinationID, source.Mode, binding.Ref, source.Limits, s.cfg.Now()) created.Placement = canonicalPlacementMetadata(binding) + brokerKeys, brokerKeysPresent := source.WorkspaceEnrollmentBrokerKeys() authority, bound := source.BoundAuthority() if !bound { authority = session.Authority{} } - if err := setSessionLabels(created, selector, profileForSession(source), source.Owner, authority); err != nil { + if err := setSessionLabels(created, selector, profileForSession(source), source.Owner, authority, brokerKeys, brokerKeysPresent); err != nil { return "", err } created.ExecutionTemplateID, created.ExecutionTemplateRevision = source.ExecutionTemplateID, source.ExecutionTemplateRevision @@ -281,7 +282,7 @@ func (s *Service) createPlacedSuccessorLocked(ctx context.Context, req ForkSucce return "", fmt.Errorf("%w: per-session engine not supported (no session-engine factory configured)", ErrInvalidArgument) } if broker != nil { - builtEngine, err = s.buildAndRegisterSessionEngineWithBrokerTools(mutationCtx, created, selector, profile, created.Mode, false, brokerTools(broker), true, inheritedSource) + builtEngine, err = s.buildAndRegisterSessionEngineWithBrokerTools(mutationCtx, created, selector, profile, created.Mode, false, brokerTools(broker), nil, true, inheritedSource) } else { builtEngine, err = s.buildAndRegisterSessionEngine(mutationCtx, created, selector, profile, created.Mode, false, inheritedSource) } diff --git a/internal/adapter/server/placement_successor_broker_test.go b/internal/adapter/server/placement_successor_broker_test.go index 572df73c8d..d1c1ed116a 100644 --- a/internal/adapter/server/placement_successor_broker_test.go +++ b/internal/adapter/server/placement_successor_broker_test.go @@ -6,6 +6,7 @@ import ( "testing" "github.com/stacklok/mecatl/engine/adapter/memstore" + "github.com/stacklok/mecatl/engine/governance" "github.com/stacklok/mecatl/engine/port" "github.com/stacklok/mecatl/engine/session" "github.com/stacklok/mecatl/engine/tool" @@ -42,15 +43,20 @@ func newBrokerSuccessorService(t *testing.T, store port.SessionStore, ids ...ses return id }, MCPBroker: runtime, + RootAuthority: func(session.SessionKind) session.Authority { + return session.Authority{CapabilitySet: governance.CapabilitySet{Tools: []string{"broker-old"}}, Provenance: "test"} + }, SessionEngine: func(context.Context, ProviderSelector, []mcp.ServerConfig, SessionProfile, string, session.PermissionMode) (SessionEngineResult, error) { t.Fatal("broker successor used ordinary SessionEngine") return SessionEngineResult{}, nil }, - SessionEngineWithTools: func(_ context.Context, _ ProviderSelector, _ []mcp.ServerConfig, _ SessionProfile, _ string, _ session.PermissionMode, tools []tool.Tool) (SessionEngineResult, error) { + SessionEngineWithTools: func(_ context.Context, _ ProviderSelector, _ []mcp.ServerConfig, _ SessionProfile, _ string, _ session.PermissionMode, tools []tool.Tool, _ []string) (SessionEngineResult, error) { if tools == nil { t.Fatal("broker successor received nil broker tools") } - return brokerEngineResult(), nil + res := brokerEngineResult() + res.BrokerRegistrationKeys = []string{"broker-old"} + return res, nil }, }) if err != nil { @@ -60,7 +66,95 @@ func newBrokerSuccessorService(t *testing.T, store port.SessionStore, ids ...ses return service } -func TestBrokerSuccessorsPersistDistinctExternalBindings(t *testing.T) { +func TestWorkspaceEnrollmentAuthority_Scenario2_PersistsAcrossResumeAndFork(t *testing.T) { + for _, tc := range []struct { + name string + clear bool + }{ + {name: "fork"}, + {name: "clear", clear: true}, + } { + t.Run(tc.name, func(t *testing.T) { + store := memstore.New() + svc := newBrokerSuccessorService(t, store, "source", "successor") + source, err := svc.CreateSession(context.Background(), session.ModeDefault, session.Limits{}) + if err != nil { + t.Fatal(err) + } + if err := store.Save(context.Background(), source); err != nil { + t.Fatal(err) + } + var successorID session.SessionID + if tc.clear { + successorID, err = svc.ClearSessionSuccessor(context.Background(), source.ID, SuccessorPlacement{}) + } else { + successorID, err = svc.ForkSessionSuccessor(context.Background(), ForkSuccessorRequest{Source: source.ID}) + } + if err != nil { + t.Fatalf("successor: %v", err) + } + successor, err := store.Load(context.Background(), successorID) + if err != nil { + t.Fatal(err) + } + keys, present := successor.WorkspaceEnrollmentBrokerKeys() + if !present || len(keys) != 1 || keys[0] != "broker-old" { + t.Fatalf("successor ledger = %q, %t", keys, present) + } + }) + } +} + +func TestBrokerSuccessorsRejectPriorBrokerKeyOwnedByNonBroker(t *testing.T) { + for _, tc := range []struct { + name string + clear bool + }{ + {name: "fork"}, + {name: "clear", clear: true}, + } { + t.Run(tc.name, func(t *testing.T) { + store := memstore.New() + svc := newBrokerSuccessorService(t, store, "source", "successor") + source, err := svc.CreateSession(context.Background(), session.ModeDefault, session.Limits{}) + if err != nil { + t.Fatal(err) + } + if err := store.Save(context.Background(), source); err != nil { + t.Fatal(err) + } + var closed bool + svc.cfg.SessionEngineWithTools = func(context.Context, ProviderSelector, []mcp.ServerConfig, SessionProfile, string, session.PermissionMode, []tool.Tool, []string) (SessionEngineResult, error) { + return SessionEngineResult{ + Engine: brokerEngineResult().Engine, + NonBrokerRegistrationKeys: []string{"broker-old"}, + Close: func() error { closed = true; return nil }, + }, nil + } + + if tc.clear { + _, err = svc.ClearSessionSuccessor(context.Background(), source.ID, SuccessorPlacement{}) + } else { + _, err = svc.ForkSessionSuccessor(context.Background(), ForkSuccessorRequest{Source: source.ID}) + } + var collision *WorkspaceEnrollmentCollisionError + if !errors.As(err, &collision) || collision.Key != "broker-old" { + t.Fatalf("successor collision = %v, want safe prior registration key", err) + } + if !closed { + t.Fatal("successor left the unpublished conflicting candidate open") + } + svc.mu.Lock() + published := svc.sessionEngines["successor"] + svc.mu.Unlock() + if published != nil { + t.Fatal("successor published a non-broker implementation under a prior broker key") + } + }) + } +} + +func TestWorkspaceEnrollmentAuthority_Scenario2_SuccessorsStartFreshBrokerEnrollment(t *testing.T) { for _, tc := range []struct { name string clear bool diff --git a/internal/adapter/server/service.go b/internal/adapter/server/service.go index 5651a05994..690a152b5e 100644 --- a/internal/adapter/server/service.go +++ b/internal/adapter/server/service.go @@ -172,10 +172,41 @@ type SessionEngineResult struct { // publication. Engines do not hold runtime pins; run admission compares this // tag with the operation pin and rebuilds before use. RuntimeRevision uint64 + // BrokerRegistrationKeys are the exact registration keys successfully mounted + // from the candidate broker bundle. NonBrokerRegistrationKeys is the complete + // final-catalog partition excluding that bundle. Both are Catalog.Names metadata. + BrokerRegistrationKeys []string + NonBrokerRegistrationKeys []string // Close tears down the session's MCP manager. Never nil (a no-op when no specs). Close func() error } +func workspaceEnrollmentRegistrationOverlap(priorBrokerKeys, authorityTools, nonBrokerRegistrationKeys []string) (string, bool) { + for _, key := range nonBrokerRegistrationKeys { + if slices.Contains(priorBrokerKeys, key) && slices.Contains(authorityTools, key) { + return key, true + } + } + return "", false +} + +func sameWorkspaceEnrollmentToolKeys(left, right []string) bool { + if len(left) != len(right) { + return false + } + keys := make(map[string]struct{}, len(left)) + for _, key := range left { + keys[key] = struct{}{} + } + for _, key := range right { + if _, ok := keys[key]; !ok { + return false + } + delete(keys, key) + } + return len(keys) == 0 +} + // ResolvedModel is the per-session EFFECTIVE model echoed on the wire: the // provider+model id this session resolved to plus its context window. It is the // SINGLE composition-computed value (see Config.DefaultResolvedModel and the @@ -250,10 +281,23 @@ type ModelInventory interface { CurrentModelSnapshot() ModelSnapshot } +// WorkspaceEnrollmentCollisionError is the presentation-safe collision result +// exchanged only between root-internal composition and server seams. +type WorkspaceEnrollmentCollisionError struct{ Key string } + +func (e *WorkspaceEnrollmentCollisionError) Error() string { + if e == nil || e.Key == "" { + return "workspace enrollment registration collision; correct the broker configuration and retry" + } + return fmt.Sprintf("workspace tool %q conflicts with an existing registration; correct the broker configuration and retry", e.Key) +} + +func (*WorkspaceEnrollmentCollisionError) Unwrap() error { return ErrFailedPrecondition } + // SessionEngineWithToolsFactory is the explicit per-session catalogue seam used // for host-owned wrapper tools. Tools are ordinary arguments: catalogue assembly // must not recover them from context values or a process-global fallback. -type SessionEngineWithToolsFactory func(ctx context.Context, sel ProviderSelector, specs []mcp.ServerConfig, profile SessionProfile, workspace string, mode session.PermissionMode, sessionTools []tool.Tool) (SessionEngineResult, error) +type SessionEngineWithToolsFactory func(ctx context.Context, sel ProviderSelector, specs []mcp.ServerConfig, profile SessionProfile, workspace string, mode session.PermissionMode, sessionTools []tool.Tool, sessionToolKeys []string) (SessionEngineResult, error) // ExecutionWorkspaceAcquirer lazily borrows the exact server-authorized session // workspace as a source-only capability. The caller owns the returned release. @@ -262,7 +306,9 @@ type ExecutionWorkspaceAcquirer func(context.Context) (tool.Workspace, func() er // SessionContextEngineFactory builds a session engine with its stored source identity. // The trailing source ID is set only while constructing a same-placement successor // carrying a live instruction snapshot; ordinary builds and rehydration receive empty. -type SessionContextEngineFactory func(context.Context, session.SessionID, *session.Principal, ExecutionWorkspaceAcquirer, ProviderSelector, []mcp.ServerConfig, SessionProfile, string, session.PermissionMode, []tool.Tool, session.SessionID) (SessionEngineResult, error) +// sessionToolKeys carries the exact frozen registration names; nil means derive +// keys from the tools, while a non-nil empty slice requests a frozen empty set. +type SessionContextEngineFactory func(context.Context, session.SessionID, *session.Principal, ExecutionWorkspaceAcquirer, ProviderSelector, []mcp.ServerConfig, SessionProfile, string, session.PermissionMode, []tool.Tool, []string, session.SessionID) (SessionEngineResult, error) // Config wires the Service's collaborators and resolved composition values. type Config struct { @@ -2081,13 +2127,23 @@ func (s *Service) createSessionWithOptions(ctx context.Context, mode session.Per // empty-selector default profile this writes the zero values, so a default // session's snapshot carries no selector/profile labels (the labels omitempty // out of the JSON). -func setSessionLabels(sess *session.Session, sel ProviderSelector, profile SessionProfile, owner *session.Principal, authority session.Authority) error { +type carriedAuthority struct { + authority session.Authority + bound bool + brokerKeys []string + brokerKeysPresent bool +} + +func setSessionLabels(sess *session.Session, sel ProviderSelector, profile SessionProfile, owner *session.Principal, authority session.Authority, brokerKeys []string, brokerKeysPresent bool) error { sess.Profile = string(profile) sess.ProviderID = sel.ProviderID sess.ModelID = sel.ModelID sess.ReasoningEffort = sel.ReasoningEffort // Owner and authority are independently stamped at the same root seam. - return sess.RestoreLabels(owner, authority) + if err := sess.RestoreLabels(owner, authority); err != nil { + return err + } + return sess.RestoreWorkspaceEnrollmentBrokerKeys(brokerKeys, brokerKeysPresent) } func (s *Service) setTitleGenerationEligibility(sess *session.Session, sel ProviderSelector) { @@ -2096,45 +2152,48 @@ func (s *Service) setTitleGenerationEligibility(sess *session.Session, sel Provi } } -func (s *Service) setPerSessionLabels(sess *session.Session, sel ProviderSelector, profile SessionProfile, owner *session.Principal, opts createSessionOpts, res SessionEngineResult, broker []tool.Tool, carried session.Authority, carriedBound bool) error { - authority := s.rootAuthority(sess.Kind, carried, carriedBound) +func (s *Service) setPerSessionLabels(sess *session.Session, sel ProviderSelector, profile SessionProfile, owner *session.Principal, opts createSessionOpts, res SessionEngineResult, carried carriedAuthority) error { + authority := s.rootAuthority(sess.Kind, carried) + brokerKeys, brokerKeysPresent := carried.brokerKeys, carried.brokerKeysPresent // Broker wrappers and client-mounted MCP tools are both resolved only after // the process root authority was minted (they are per-SESSION, not known at // build time). Include this session's exact set in a fresh root without // widening authority carried from another session. - if !carriedBound && (len(broker) != 0 || len(res.MountedClientMCPTools) != 0) { - seen := make(map[string]struct{}, len(authority.CapabilitySet.Tools)+len(broker)+len(res.MountedClientMCPTools)) - for _, name := range authority.CapabilitySet.Tools { - seen[name] = struct{}{} - } - for _, candidate := range broker { - name := candidate.Spec().Name - if _, ok := seen[name]; ok { - continue + if !carried.bound { + brokerKeys, brokerKeysPresent = res.BrokerRegistrationKeys, true + if len(res.BrokerRegistrationKeys) != 0 || len(res.MountedClientMCPTools) != 0 { + seen := make(map[string]struct{}, len(authority.CapabilitySet.Tools)+len(res.BrokerRegistrationKeys)+len(res.MountedClientMCPTools)) + for _, name := range authority.CapabilitySet.Tools { + seen[name] = struct{}{} } - seen[name] = struct{}{} - authority.CapabilitySet.Tools = append(authority.CapabilitySet.Tools, name) - } - for _, name := range res.MountedClientMCPTools { - if _, ok := seen[name]; ok { - continue + for _, name := range res.BrokerRegistrationKeys { + if _, ok := seen[name]; ok { + continue + } + seen[name] = struct{}{} + authority.CapabilitySet.Tools = append(authority.CapabilitySet.Tools, name) } - seen[name] = struct{}{} - authority.CapabilitySet.Tools = append(authority.CapabilitySet.Tools, name) + for _, name := range res.MountedClientMCPTools { + if _, ok := seen[name]; ok { + continue + } + seen[name] = struct{}{} + authority.CapabilitySet.Tools = append(authority.CapabilitySet.Tools, name) + } + sort.Strings(authority.CapabilitySet.Tools) } - sort.Strings(authority.CapabilitySet.Tools) } if sess.Kind == session.SessionKindDebug { authority.CapabilitySet.Tools = append(authority.CapabilitySet.Tools, res.DebugMCPTools...) sess.DebugMCPServers = append([]string(nil), opts.debugMCPServers...) sess.DebugMCPTools = append([]string(nil), res.DebugMCPTools...) } - return setSessionLabels(sess, sel, profile, owner, authority) + return setSessionLabels(sess, sel, profile, owner, authority, brokerKeys, brokerKeysPresent) } -func (s *Service) rootAuthority(kind session.SessionKind, carried session.Authority, carriedBound bool) session.Authority { - if carriedBound { - return carried +func (s *Service) rootAuthority(kind session.SessionKind, carried carriedAuthority) session.Authority { + if carried.bound { + return carried.authority } if s.cfg.RootAuthority == nil { return session.Authority{} @@ -2567,15 +2626,14 @@ func (s *Service) createSession(ctx context.Context, mode session.PermissionMode // A carryover fork replaces the context owner with the source's owner below. var carrySnap []session.Message - var carriedAuthority session.Authority - carriedAuthorityBound := false + var carried carriedAuthority if opts.sourceSessionID != "" { - snap, srcOwner, sourceAuthority, sourceAuthorityBound, err := s.validateCarryover(ctx, opts.sourceSessionID, sel.ProviderID) + snap, srcOwner, sourceCarried, err := s.validateCarryover(ctx, opts.sourceSessionID, sel.ProviderID) if err != nil { return nil, err } carrySnap = snap - carriedAuthority, carriedAuthorityBound = sourceAuthority, sourceAuthorityBound + carried = sourceCarried // A fork inherits the SOURCE's owner, overriding the context principal // (and any WithOwner) — see validateCarryover. owner = srcOwner @@ -2583,6 +2641,20 @@ func (s *Service) createSession(ctx context.Context, mode session.PermissionMode needPerSession := s.cfg.MCPBroker != nil || s.sessionNeedsPerFactory(sel, specs, profile, workspace) || s.cfg.LearnedSkills != nil if !needPerSession { + // Workspace broker wrappers only enter per-session catalogues; the shared + // engine can therefore identify a conflicting non-broker registration by + // exact catalog membership without forcing this create onto a factory. + if carried.brokerKeysPresent { + sharedRegistrationKeys := make([]string, 0, len(carried.brokerKeys)) + for _, key := range carried.brokerKeys { + if s.cfg.Engine.HasTool(key) { + sharedRegistrationKeys = append(sharedRegistrationKeys, key) + } + } + if key, overlaps := workspaceEnrollmentRegistrationOverlap(carried.brokerKeys, carried.authority.CapabilitySet.Tools, sharedRegistrationKeys); overlaps { + return nil, &WorkspaceEnrollmentCollisionError{Key: key} + } + } // Shared-engine fast path. The labels are // the empty pair + default profile here (the empty-selector default profile is // exactly the no-per-session case), so setLabels persists nothing new — the @@ -2591,7 +2663,11 @@ func (s *Service) createSession(ctx context.Context, mode session.PermissionMode if err != nil { return nil, fmt.Errorf("server: create session metadata: %w", err) } - if err := setSessionLabels(sess, sel, profile, owner, s.rootAuthority(sess.Kind, carriedAuthority, carriedAuthorityBound)); err != nil { + brokerKeysPresent := true + if carried.bound { + brokerKeysPresent = carried.brokerKeysPresent + } + if err := setSessionLabels(sess, sel, profile, owner, s.rootAuthority(sess.Kind, carried), carried.brokerKeys, brokerKeysPresent); err != nil { return nil, err } if opts.execution.Kind == PlacementSelectorTemplate { @@ -2611,7 +2687,7 @@ func (s *Service) createSession(ctx context.Context, mode session.PermissionMode return created, err } - return s.createPerSessionEngine(ctx, mintID, mode, workspace, limits, sel, specs, profile, carrySnap, owner, opts, carriedAuthority, carriedAuthorityBound, retryRequest, placement, &publishedPlacement) + return s.createPerSessionEngine(ctx, mintID, mode, workspace, limits, sel, specs, profile, carrySnap, owner, opts, carried, retryRequest, placement, &publishedPlacement) } // createPerSessionEngine is the per-session-engine create branch, factored out @@ -2624,7 +2700,7 @@ func (s *Service) createSession(ctx context.Context, mode session.PermissionMode // profile-aware workspace rule and the carryover snapshot semantics. // //nolint:gocyclo // Creation keeps factory, authorization, broker ownership, registration, and teardown in one transaction. -func (s *Service) createPerSessionEngine(ctx context.Context, mintID func() session.SessionID, mode session.PermissionMode, workspace string, limits session.Limits, sel ProviderSelector, specs []mcp.ServerConfig, profile SessionProfile, carrySnap []session.Message, owner *session.Principal, opts createSessionOpts, carriedAuthority session.Authority, carriedAuthorityBound bool, retryRequest *createRequest, placement *PlacementBinding, publishedPlacement *bool) (*session.Session, error) { +func (s *Service) createPerSessionEngine(ctx context.Context, mintID func() session.SessionID, mode session.PermissionMode, workspace string, limits session.Limits, sel ProviderSelector, specs []mcp.ServerConfig, profile SessionProfile, carrySnap []session.Message, owner *session.Principal, opts createSessionOpts, carried carriedAuthority, retryRequest *createRequest, placement *PlacementBinding, publishedPlacement *bool) (*session.Session, error) { if opts.debugTargetID != "" { if s.cfg.DebugSessionEngine == nil { return nil, fmt.Errorf("%w: session debugging is not supported", ErrInvalidArgument) @@ -2685,7 +2761,7 @@ func (s *Service) createPerSessionEngine(ctx context.Context, mintID func() sess } factoryDone := creatediag.Begin(ctx, "engine_factory") acquire := s.executionWorkspaceAcquirer(owner, placement.Ref) - res, err = s.callSessionEngine(ctx, id, owner, acquire, sel, specs, profile, workspace, mode, brokerTools(broker)) + res, err = s.callSessionEngine(ctx, id, owner, acquire, sel, specs, profile, workspace, mode, brokerTools(broker), nil) // Only this reserved, unpublished create may retry a retired attempt. // Ordinary rehydration must wait for explicit owner-authorized Load. if errors.Is(err, ErrCommandBindingRetired) && s.cfg.SessionContextEngine != nil { @@ -2695,7 +2771,7 @@ func (s *Service) createPerSessionEngine(ctx context.Context, mintID func() sess err = s.cfg.Commands.Activate(ctx, id, owner.Clone(), string(profile)) } if err == nil { - res, err = s.callSessionEngine(ctx, id, owner, acquire, sel, specs, profile, workspace, mode, brokerTools(broker)) + res, err = s.callSessionEngine(ctx, id, owner, acquire, sel, specs, profile, workspace, mode, brokerTools(broker), nil) } } factoryDone(err) @@ -2705,6 +2781,14 @@ func (s *Service) createPerSessionEngine(ctx context.Context, mintID func() sess // error is propagated as-is for the caller to map to a status. return nil, err } + if carried.brokerKeysPresent { + if key, overlaps := workspaceEnrollmentRegistrationOverlap(carried.brokerKeys, carried.authority.CapabilitySet.Tools, res.NonBrokerRegistrationKeys); overlaps { + if res.Close != nil { + _ = res.Close() + } + return nil, &WorkspaceEnrollmentCollisionError{Key: key} + } + } eng, closeFn := res.Engine, res.Close // All-or-nothing client MCP on the wire path, before a non-broker id is minted // or anything is persisted: a caller that asked for tools must not be handed a @@ -2734,7 +2818,7 @@ func (s *Service) createPerSessionEngine(ctx context.Context, mintID func() sess // creation labels on the aggregate, so a restarted process re-derives the SAME // per-session engine via the factory (rehydrateSession) instead of falling to the // default-provider floor / inferring the profile from the empty-workspace pun. - if err := s.setPerSessionLabels(sess, sel, profile, owner, opts, res, brokerTools(broker), carriedAuthority, carriedAuthorityBound); err != nil { + if err := s.setPerSessionLabels(sess, sel, profile, owner, opts, res, carried); err != nil { if closeFn != nil { _ = closeFn() } @@ -4464,13 +4548,13 @@ func (s *Service) LoadSession(ctx context.Context, id session.SessionID) (*sessi // default). The model is intentionally NOT compared: a model mismatch within a // provider is the point of the /models picker switch, so a same-provider // model change is permitted; a cross-provider model change strips and replays. -func (s *Service) validateCarryover(ctx context.Context, srcID session.SessionID, newProviderID string) ([]session.Message, *session.Principal, session.Authority, bool, error) { +func (s *Service) validateCarryover(ctx context.Context, srcID session.SessionID, newProviderID string) ([]session.Message, *session.Principal, carriedAuthority, error) { src, err := s.loadAndReopen(ctx, srcID) if err != nil { - return nil, nil, session.Authority{}, false, err + return nil, nil, carriedAuthority{}, err } if src.State == session.StateRunning || src.State == session.StateAwaiting || src.State == session.StateAuthorizing { - return nil, nil, session.Authority{}, false, fmt.Errorf("%w: carryover requires a session at a turn boundary; source %q is %s", ErrFailedPrecondition, srcID, src.State) + return nil, nil, carriedAuthority{}, fmt.Errorf("%w: carryover requires a session at a turn boundary; source %q is %s", ErrFailedPrecondition, srcID, src.State) } // The SOURCE's owner travels with the carried history: // a fork is attributed to whoever owned the session it copied, never to the @@ -4479,7 +4563,13 @@ func (s *Service) validateCarryover(ctx context.Context, srcID session.SessionID // yields an ownerless fork, never a fabricated one. srcOwner := src.Owner srcAuthority, srcAuthorityBound := src.BoundAuthority() - return s.providerCarryoverSnapshot(src, newProviderID), srcOwner, srcAuthority, srcAuthorityBound, nil + brokerKeys, brokerKeysPresent := src.WorkspaceEnrollmentBrokerKeys() + return s.providerCarryoverSnapshot(src, newProviderID), srcOwner, carriedAuthority{ + authority: srcAuthority, + bound: srcAuthorityBound, + brokerKeys: brokerKeys, + brokerKeysPresent: brokerKeysPresent, + }, nil } func (s *Service) providerCarryoverSnapshot(src *session.Session, newProviderID string) []session.Message { @@ -4851,7 +4941,7 @@ func (s *Service) LoadSessionWithMCP(ctx context.Context, id session.SessionID, if err != nil { return nil, err } - res, err := s.callSessionEngine(ctx, sess.ID, sess.Owner.Clone(), s.executionWorkspaceAcquirer(sess.Owner, sess.EnvironmentRef), sel, specs, profile, workspace, sess.Mode, nil) + res, err := s.callSessionEngine(ctx, sess.ID, sess.Owner.Clone(), s.executionWorkspaceAcquirer(sess.Owner, sess.EnvironmentRef), sel, specs, profile, workspace, sess.Mode, nil, nil) if err != nil { // The session was loaded + (if needed) reopened and re-persisted, but the // per-session engine could not be built. We deliberately do NOT roll that @@ -4861,6 +4951,14 @@ func (s *Service) LoadSessionWithMCP(ctx context.Context, id session.SessionID, // failure tears the freshly-built engine down). return nil, err } + if priorBrokerKeys, present := sess.WorkspaceEnrollmentBrokerKeys(); present { + if key, overlaps := workspaceEnrollmentRegistrationOverlap(priorBrokerKeys, sess.Authority.CapabilitySet.Tools, res.NonBrokerRegistrationKeys); overlaps { + if res.Close != nil { + _ = res.Close() + } + return nil, &WorkspaceEnrollmentCollisionError{Key: key} + } + } s.mu.Lock() // Re-load leak guard: if a per-session engine is already registered for this id // (e.g. a re-load of the same session on the same connection), close the prior @@ -6378,7 +6476,7 @@ func (s *Service) rehydrateSession(ctx context.Context, sess *session.Session) ( // upper-bound capability record until the owner explicitly starts a new // enrollment. Exact-tool mode with no tools keeps ordinary prompts usable while // making persisted broker names non-executable. - return s.buildAndRegisterSessionEngineWithBrokerTools(ctx, sess, sel, profile, sess.Mode, false, nil, true) + return s.buildAndRegisterSessionEngineWithBrokerTools(ctx, sess, sel, profile, sess.Mode, false, nil, nil, true) } // buildAndRegisterSessionEngine is the ONE shared build+cap-check+register+teardown @@ -6406,11 +6504,11 @@ func (s *Service) rehydrateSession(ctx context.Context, sess *session.Session) ( // //nolint:gocyclo // Explicit validation, rebuild, broker, capacity, and rollback gates stay ordered. func (s *Service) buildAndRegisterSessionEngine(ctx context.Context, sess *session.Session, sel ProviderSelector, profile SessionProfile, mode session.PermissionMode, replace bool, forkSource ...session.SessionID) (*sessionEngine, error) { - return s.buildAndRegisterSessionEngineWithBrokerTools(ctx, sess, sel, profile, mode, replace, nil, false, forkSource...) + return s.buildAndRegisterSessionEngineWithBrokerTools(ctx, sess, sel, profile, mode, replace, nil, nil, false, forkSource...) } //nolint:gocyclo // rehydration keeps validation, factory selection, broker, capacity, and rollback gates ordered; inherent. -func (s *Service) buildAndRegisterSessionEngineWithBrokerTools(ctx context.Context, sess *session.Session, sel ProviderSelector, profile SessionProfile, mode session.PermissionMode, replace bool, exactTools []tool.Tool, useExactTools bool, forkSource ...session.SessionID) (*sessionEngine, error) { +func (s *Service) buildAndRegisterSessionEngineWithBrokerTools(ctx context.Context, sess *session.Session, sel ProviderSelector, profile SessionProfile, mode session.PermissionMode, replace bool, exactTools []tool.Tool, exactToolNames []string, useExactTools bool, forkSource ...session.SessionID) (*sessionEngine, error) { id := sess.ID unlockBroker := s.brokerMu.lock(id) defer unlockBroker() @@ -6424,6 +6522,11 @@ func (s *Service) buildAndRegisterSessionEngineWithBrokerTools(ctx context.Conte return nil, fmt.Errorf("%w: %d", ErrTooManySessionEngines, s.cfg.MaxSessionEngines) } } + if useExactTools && exactToolNames != nil { + if _, present := sess.WorkspaceEnrollmentBrokerKeys(); !present { + return nil, fmt.Errorf("%w: workspace enrollment provenance is unavailable", ErrFailedPrecondition) + } + } var res SessionEngineResult var err error var broker *localBrokerAttachment @@ -6449,7 +6552,7 @@ func (s *Service) buildAndRegisterSessionEngineWithBrokerTools(ctx context.Conte if workspaceErr != nil { return nil, workspaceErr } - res, err = s.callSessionEngine(ctx, sess.ID, sess.Owner.Clone(), s.executionWorkspaceAcquirer(sess.Owner, sess.EnvironmentRef), sel, specs, profile, workspace, mode, append([]tool.Tool(nil), exactTools...), forkSource...) + res, err = s.callSessionEngine(ctx, sess.ID, sess.Owner.Clone(), s.executionWorkspaceAcquirer(sess.Owner, sess.EnvironmentRef), sel, specs, profile, workspace, mode, append([]tool.Tool(nil), exactTools...), exactToolNames, forkSource...) } else { broker, err = s.openBrokerAttachment(ctx, id, sess.ExternalBinding, true) if err != nil { @@ -6460,11 +6563,28 @@ func (s *Service) buildAndRegisterSessionEngineWithBrokerTools(ctx context.Conte if workspaceErr != nil { return nil, workspaceErr } - res, err = s.callSessionEngine(ctx, sess.ID, sess.Owner.Clone(), s.executionWorkspaceAcquirer(sess.Owner, sess.EnvironmentRef), sel, specs, profile, workspace, mode, brokerTools(broker), forkSource...) + res, err = s.callSessionEngine(ctx, sess.ID, sess.Owner.Clone(), s.executionWorkspaceAcquirer(sess.Owner, sess.EnvironmentRef), sel, specs, profile, workspace, mode, brokerTools(broker), nil, forkSource...) } if err != nil { return nil, fmt.Errorf("server: build session engine %q: %w", id, err) } + closeCandidate := func() { + if res.Close != nil { + _ = res.Close() + } + } + if useExactTools && exactToolNames != nil { + if !sameWorkspaceEnrollmentToolKeys(exactToolNames, res.BrokerRegistrationKeys) { + closeCandidate() + return nil, fmt.Errorf("%w: workspace enrollment catalogue registration metadata is inconsistent", ErrFailedPrecondition) + } + } + if priorBrokerKeys, present := sess.WorkspaceEnrollmentBrokerKeys(); present { + if key, overlaps := workspaceEnrollmentRegistrationOverlap(priorBrokerKeys, sess.Authority.CapabilitySet.Tools, res.NonBrokerRegistrationKeys); overlaps { + closeCandidate() + return nil, &WorkspaceEnrollmentCollisionError{Key: key} + } + } se := &sessionEngine{ engine: res.Engine, caps: res.Capabilities, @@ -6481,16 +6601,12 @@ func (s *Service) buildAndRegisterSessionEngineWithBrokerTools(ctx context.Conte // FIRST-REGISTRATION-WINS: a concurrent rehydration (or load) won the race; // keep its engine, tear ours down so the loser leaks no MCP manager. s.mu.Unlock() - if se.close != nil { - _ = se.close() - } + closeCandidate() return prior, nil } if !replace && len(s.sessionEngines) >= s.cfg.MaxSessionEngines { s.mu.Unlock() - if se.close != nil { - _ = se.close() - } + closeCandidate() return nil, fmt.Errorf("%w: %d", ErrTooManySessionEngines, s.cfg.MaxSessionEngines) } if replace { diff --git a/internal/adapter/server/workspace_enrollment.go b/internal/adapter/server/workspace_enrollment.go index 198e82e6d5..9050750a44 100644 --- a/internal/adapter/server/workspace_enrollment.go +++ b/internal/adapter/server/workspace_enrollment.go @@ -44,26 +44,7 @@ func (s *Service) connectWorkspaceServicesLocked(ctx context.Context, id session pending, exists := sess.PendingWorkspaceEnrollment() if !exists { - if len(sess.Conversation.Messages) != 0 { - if err := enroller.ResetWorkspaceEnrollment(ctx); err != nil { - return WorkspaceEnrollmentProjection{}, fmt.Errorf("%w: reset workspace enrollment", ErrFailedPrecondition) - } - s.withdrawBrokerEngine(id) - } - presentation, beginErr := enroller.BeginWorkspaceEnrollment(ctx) - if beginErr != nil || !presentation.Valid() { - return WorkspaceEnrollmentProjection{}, fmt.Errorf("%w: begin workspace enrollment", ErrFailedPrecondition) - } - pending = pendingEnrollment(presentation.Ref) - if err := sess.BeginWorkspaceEnrollment(pending); err != nil { - cancelWorkspaceEnrollmentDetached(ctx, enroller, presentation.Ref) - return WorkspaceEnrollmentProjection{}, fmt.Errorf("%w: record workspace enrollment", ErrFailedPrecondition) - } - if err := s.saveSession(ctx, sess); err != nil { - cancelWorkspaceEnrollmentDetached(ctx, enroller, presentation.Ref) - return WorkspaceEnrollmentProjection{}, fmt.Errorf("%w: persist workspace enrollment", ErrInternal) - } - return WorkspaceEnrollmentProjection{Ref: presentation.Ref, Status: brokercontract.WorkspaceEnrollmentPending, URL: presentation.URL}, nil + return s.beginWorkspaceEnrollment(ctx, id, sess, enroller) } expectedRef := enrollmentRef(pending) @@ -85,12 +66,23 @@ func (s *Service) connectWorkspaceServicesLocked(ctx context.Context, id session release() release = func() {} sel := ProviderSelector{ProviderID: sess.ProviderID, ModelID: sess.ModelID, ReasoningEffort: sess.ReasoningEffort} - if _, err := s.buildAndRegisterSessionEngineWithBrokerTools(ctx, sess, sel, profileForSession(sess), sess.Mode, true, exactTools, true); err != nil { + _, err = s.buildAndRegisterSessionEngineWithBrokerTools(ctx, sess, sel, profileForSession(sess), sess.Mode, true, exactTools, toolNames, true) + if err != nil { s.withdrawBrokerEngine(sess.ID) + var collision *WorkspaceEnrollmentCollisionError + if errors.As(err, &collision) { + if _, settleErr := s.settleTerminalWorkspaceEnrollment(ctx, sess, brokercontract.WorkspaceEnrollmentFailed); settleErr != nil { + return WorkspaceEnrollmentProjection{}, settleErr + } + return WorkspaceEnrollmentProjection{}, collision + } return WorkspaceEnrollmentProjection{}, err } if err := sess.CompleteWorkspaceEnrollment(pending, toolNames); err != nil { s.withdrawBrokerEngine(sess.ID) + if _, settleErr := s.settleTerminalWorkspaceEnrollment(ctx, sess, brokercontract.WorkspaceEnrollmentFailed); settleErr != nil { + return WorkspaceEnrollmentProjection{}, settleErr + } return WorkspaceEnrollmentProjection{}, fmt.Errorf("%w: complete workspace enrollment", ErrFailedPrecondition) } if err := s.saveSession(ctx, sess); err != nil { @@ -100,6 +92,29 @@ func (s *Service) connectWorkspaceServicesLocked(ctx context.Context, id session return WorkspaceEnrollmentProjection{Ref: result.Ref, Status: result.Status}, nil } +func (s *Service) beginWorkspaceEnrollment(ctx context.Context, id session.SessionID, sess *session.Session, enroller brokercontract.WorkspaceEnrollmentAttachment) (WorkspaceEnrollmentProjection, error) { + if len(sess.Conversation.Messages) != 0 { + if err := enroller.ResetWorkspaceEnrollment(ctx); err != nil { + return WorkspaceEnrollmentProjection{}, fmt.Errorf("%w: reset workspace enrollment", ErrFailedPrecondition) + } + s.withdrawBrokerEngine(id) + } + presentation, err := enroller.BeginWorkspaceEnrollment(ctx) + if err != nil || !presentation.Valid() { + return WorkspaceEnrollmentProjection{}, fmt.Errorf("%w: begin workspace enrollment", ErrFailedPrecondition) + } + pending := pendingEnrollment(presentation.Ref) + if err := sess.BeginWorkspaceEnrollment(pending); err != nil { + cancelWorkspaceEnrollmentDetached(ctx, enroller, presentation.Ref) + return WorkspaceEnrollmentProjection{}, fmt.Errorf("%w: record workspace enrollment", ErrFailedPrecondition) + } + if err := s.saveSession(ctx, sess); err != nil { + cancelWorkspaceEnrollmentDetached(ctx, enroller, presentation.Ref) + return WorkspaceEnrollmentProjection{}, fmt.Errorf("%w: persist workspace enrollment", ErrInternal) + } + return WorkspaceEnrollmentProjection{Ref: presentation.Ref, Status: brokercontract.WorkspaceEnrollmentPending, URL: presentation.URL}, nil +} + func (s *Service) recordObservedWorkspaceEnrollment(ctx context.Context, sess *session.Session, pending session.PendingWorkspaceEnrollment, result brokercontract.WorkspaceEnrollmentResult) (WorkspaceEnrollmentProjection, error) { if result.Status != brokercontract.WorkspaceEnrollmentPending { if err := sess.AbortWorkspaceEnrollment(pending.ID); err != nil { @@ -207,6 +222,9 @@ func (s *Service) workspaceEnrollmentTarget(ctx context.Context, id session.Sess if err != nil || sess == nil || sess.ID != id || s.authorizeSession(ctx, sess) != nil { return nil, nil, nil, ErrNotFound } + if _, present := sess.WorkspaceEnrollmentBrokerKeys(); !present { + return sess, nil, nil, fmt.Errorf("%w: workspace enrollment provenance is unavailable; create a fresh session", ErrFailedPrecondition) + } if sess.State == session.StateCompleted { if err := sess.Reopen(); err != nil { return sess, nil, nil, fmt.Errorf("%w: reopen workspace enrollment session", ErrFailedPrecondition) diff --git a/internal/adapter/server/workspace_enrollment_test.go b/internal/adapter/server/workspace_enrollment_test.go index b14d27ac26..d176349e3c 100644 --- a/internal/adapter/server/workspace_enrollment_test.go +++ b/internal/adapter/server/workspace_enrollment_test.go @@ -4,11 +4,13 @@ import ( "context" "encoding/json" "errors" + "fmt" "go/ast" "go/parser" "go/token" "net/http" "net/http/httptest" + "reflect" "slices" "strings" "sync" @@ -19,7 +21,10 @@ import ( mecatlv1 "github.com/stacklok/mecatl/contracts/gen/go/mecatl/v1" "github.com/stacklok/mecatl/engine/adapter/memlease" "github.com/stacklok/mecatl/engine/adapter/memstore" + "github.com/stacklok/mecatl/engine/adapter/mockllm" + "github.com/stacklok/mecatl/engine/adapter/permpolicy" "github.com/stacklok/mecatl/engine/adapter/wallclock" + "github.com/stacklok/mecatl/engine/agent" "github.com/stacklok/mecatl/engine/governance" "github.com/stacklok/mecatl/engine/port" "github.com/stacklok/mecatl/engine/session" @@ -119,13 +124,19 @@ func (b *enrollmentBroker) AttachSession(ctx context.Context, id session.Session return b.attachment, outcome, nil } -type enrollmentTool struct{ name string } +type enrollmentTool struct { + name string + executed *atomic.Int32 +} func (t enrollmentTool) Spec() tool.ToolSpec { return tool.ToolSpec{Name: t.name, Schema: json.RawMessage(`{"type":"object"}`)} } func (enrollmentTool) ReadOnly() bool { return true } -func (enrollmentTool) Execute(context.Context, session.ToolCall, tool.Environment) (session.ToolResult, error) { +func (t enrollmentTool) Execute(context.Context, session.ToolCall, tool.Environment) (session.ToolResult, error) { + if t.executed != nil { + t.executed.Add(1) + } return session.NewToolResult("call", "ok"), nil } @@ -147,8 +158,8 @@ func newWorkspaceEnrollmentServiceWithStore(t *testing.T, ownershipEnforced bool RootAuthority: func(session.SessionKind) session.Authority { return session.Authority{CapabilitySet: governance.CapabilitySet{Tools: []string{"mcp__calendar__list"}}, Provenance: "test"} }, - SessionEngineWithTools: func(context.Context, ProviderSelector, []mcp.ServerConfig, SessionProfile, string, session.PermissionMode, []tool.Tool) (SessionEngineResult, error) { - return brokerEngineResult(), nil + SessionEngineWithTools: func(_ context.Context, _ ProviderSelector, _ []mcp.ServerConfig, _ SessionProfile, _ string, _ session.PermissionMode, tools []tool.Tool, _ []string) (SessionEngineResult, error) { + return brokerEngineResultWithTools(tools), nil }, }) if err != nil { @@ -216,11 +227,11 @@ func newDestructiveReplacementFixtureWithStore(t *testing.T, store port.SessionS RootAuthority: func(session.SessionKind) session.Authority { return session.Authority{CapabilitySet: governance.CapabilitySet{Tools: []string{"mcp__calendar__list"}}, Provenance: "test"} }, - SessionEngineWithTools: func(context.Context, ProviderSelector, []mcp.ServerConfig, SessionProfile, string, session.PermissionMode, []tool.Tool) (SessionEngineResult, error) { + SessionEngineWithTools: func(_ context.Context, _ ProviderSelector, _ []mcp.ServerConfig, _ SessionProfile, _ string, _ session.PermissionMode, tools []tool.Tool, _ []string) (SessionEngineResult, error) { if f.failEngineBuild.Load() { return SessionEngineResult{}, errFixtureEngineBuildFailed } - return brokerEngineResultWithStore(store), nil + return brokerEngineResultWithStoreAndTools(store, tools), nil }, }) if err != nil { @@ -258,6 +269,169 @@ func newDestructiveReplacementFixtureWithStore(t *testing.T, store port.SessionS return f } +func TestWorkspaceEnrollmentAuthority_Scenario1_CreationRecordsBrokerKeys(t *testing.T) { + store := memstore.New() + runtime := testBrokerRuntime(t) + t.Cleanup(func() { _ = runtime.Close() }) + broker := &enrollmentBroker{Service: runtime} + svc, err := NewService(Config{ + Engine: brokerEngineResult().Engine, + Store: store, + PlacementProvider: brokerPlacementProvider{}, + PlacementScope: "test", + NewID: func() session.SessionID { return "creation-records-broker-keys" }, + MCPBroker: broker, + RootAuthority: func(session.SessionKind) session.Authority { + return session.Authority{CapabilitySet: governance.CapabilitySet{Tools: []string{"Read"}}, Provenance: "test"} + }, + SessionEngineWithTools: func(_ context.Context, _ ProviderSelector, _ []mcp.ServerConfig, _ SessionProfile, _ string, _ session.PermissionMode, tools []tool.Tool, _ []string) (SessionEngineResult, error) { + return brokerEngineResultWithTools(tools), nil + }, + }) + if err != nil { + t.Fatal(err) + } + t.Cleanup(svc.Close) + + created, err := svc.CreateSession(t.Context(), session.ModeDefault, session.Limits{}) + if err != nil { + t.Fatal(err) + } + stored, err := store.Load(t.Context(), created.ID) + if err != nil { + t.Fatal(err) + } + if keys, present := stored.WorkspaceEnrollmentBrokerKeys(); !present || !slices.Equal(keys, []string{"mcp__calendar__list"}) { + t.Fatalf("created broker ledger = (%q, %t), want ([mcp__calendar__list], true)", keys, present) + } + + started, err := svc.ConnectWorkspaceServices(t.Context(), created.ID) + if err != nil { + t.Fatal(err) + } + catalogue, err := brokercontract.NewWorkspaceCatalogue(started.Ref, []tool.Tool{enrollmentTool{name: "mcp__slack__post"}}) + if err != nil { + t.Fatal(err) + } + broker.attachment.result = brokercontract.WorkspaceEnrollmentResult{Ref: started.Ref, Status: brokercontract.WorkspaceEnrollmentConnected, Catalogue: catalogue} + if _, err := svc.ConnectWorkspaceServices(t.Context(), created.ID); err != nil { + t.Fatal(err) + } + stored, err = store.Load(t.Context(), created.ID) + if err != nil { + t.Fatal(err) + } + authority, bound := stored.BoundAuthority() + if !bound || !slices.Equal(authority.CapabilitySet.Tools, []string{"Read", "mcp__slack__post"}) { + t.Fatalf("authority after enrollment = (%q, %t), want ([Read mcp__slack__post], true)", authority.CapabilitySet.Tools, bound) + } +} + +func TestWorkspaceEnrollmentAuthority_Scenario2_CarryoverCopiesProvenance(t *testing.T) { + store := memstore.New() + runtime := testBrokerRuntime(t) + t.Cleanup(func() { _ = runtime.Close() }) + broker := &enrollmentBroker{Service: runtime} + var next int + svc, err := NewService(Config{ + Engine: brokerEngineResult().Engine, + Store: store, + PlacementProvider: brokerPlacementProvider{}, + PlacementScope: "test", + NewID: func() session.SessionID { + next++ + return session.SessionID(fmt.Sprintf("carryover-%d", next)) + }, + MCPBroker: broker, + RootAuthority: func(session.SessionKind) session.Authority { + return session.Authority{CapabilitySet: governance.CapabilitySet{Tools: []string{"Read"}}, Provenance: "test"} + }, + SessionEngineWithTools: func(_ context.Context, _ ProviderSelector, _ []mcp.ServerConfig, _ SessionProfile, _ string, _ session.PermissionMode, tools []tool.Tool, _ []string) (SessionEngineResult, error) { + return brokerEngineResultWithTools(tools), nil + }, + }) + if err != nil { + t.Fatal(err) + } + t.Cleanup(svc.Close) + + makeSource := func(t *testing.T, id session.SessionID, keys []string, present bool) *session.Session { + t.Helper() + source := session.New(id, session.ModeDefault, session.EnvironmentRef{Kind: session.EnvKindLocal, ID: "/workspace", Revision: "in-tree-v1"}, session.Limits{}, time.Now()) + if err := source.BindAuthority(session.Authority{CapabilitySet: governance.CapabilitySet{Tools: []string{"Read", "mcp__calendar__list"}}, Provenance: "test"}); err != nil { + t.Fatal(err) + } + if present { + if err := source.RestoreWorkspaceEnrollmentBrokerKeys(keys, true); err != nil { + t.Fatal(err) + } + } + if err := store.Create(t.Context(), source); err != nil { + t.Fatal(err) + } + return source + } + + t.Run("present", func(t *testing.T) { + source := makeSource(t, "carryover-present", []string{"mcp__calendar__list"}, true) + created, err := svc.CreateSessionWithProfile(t.Context(), session.ModeDefault, session.Limits{}, ProviderSelector{}, ProfileDefault, WithSourceSession(source.ID)) + if err != nil { + t.Fatal(err) + } + stored, err := store.Load(t.Context(), created.ID) + if err != nil { + t.Fatal(err) + } + if keys, present := stored.WorkspaceEnrollmentBrokerKeys(); !present || !slices.Equal(keys, []string{"mcp__calendar__list"}) { + t.Fatalf("carried broker ledger = (%q, %t), want ([mcp__calendar__list], true)", keys, present) + } + }) + + t.Run("absent legacy is retained and enrollment is refused", func(t *testing.T) { + source := makeSource(t, "carryover-absent", nil, false) + created, err := svc.CreateSessionWithProfile(t.Context(), session.ModeDefault, session.Limits{}, ProviderSelector{}, ProfileDefault, WithSourceSession(source.ID)) + if err != nil { + t.Fatal(err) + } + stored, err := store.Load(t.Context(), created.ID) + if err != nil { + t.Fatal(err) + } + if keys, present := stored.WorkspaceEnrollmentBrokerKeys(); present || len(keys) != 0 { + t.Fatalf("legacy carried broker ledger = (%q, %t), want ([], false)", keys, present) + } + if _, err := svc.ConnectWorkspaceServices(t.Context(), created.ID); !errors.Is(err, ErrFailedPrecondition) { + t.Fatalf("enrollment error = %v, want ErrFailedPrecondition", err) + } + }) +} + +func replacePersistedBrokerLedger(t *testing.T, svc *Service, created *session.Session, keys []string, present bool) { + t.Helper() + legacy := session.New(created.ID, created.Mode, created.EnvironmentRef, created.Limits, created.CreatedAt) + legacy.ExternalBinding = created.ExternalBinding + authority, bound := created.BoundAuthority() + if bound && present { + for _, key := range keys { + if !slices.Contains(authority.CapabilitySet.Tools, key) { + authority.CapabilitySet.Tools = append(authority.CapabilitySet.Tools, key) + } + } + slices.Sort(authority.CapabilitySet.Tools) + } + if err := legacy.RestoreLabels(created.Owner, authority); err != nil { + t.Fatal(err) + } + if present { + if err := legacy.RestoreWorkspaceEnrollmentBrokerKeys(keys, true); err != nil { + t.Fatal(err) + } + } + if err := svc.cfg.Store.Save(t.Context(), legacy); err != nil { + t.Fatal(err) + } +} + func (f *destructiveReplacementFixture) authorityTools(t *testing.T) []string { t.Helper() loaded, err := f.store.Load(t.Context(), f.created.ID) @@ -301,8 +475,8 @@ func TestIdleSessionBrokerRefresh_Scenario1_ExistingConversationControls(t *test RootAuthority: func(session.SessionKind) session.Authority { return session.Authority{CapabilitySet: governance.CapabilitySet{Tools: []string{"mcp__calendar__list"}}, Provenance: "test"} }, - SessionEngineWithTools: func(context.Context, ProviderSelector, []mcp.ServerConfig, SessionProfile, string, session.PermissionMode, []tool.Tool) (SessionEngineResult, error) { - return brokerEngineResultWithStore(store), nil + SessionEngineWithTools: func(_ context.Context, _ ProviderSelector, _ []mcp.ServerConfig, _ SessionProfile, _ string, _ session.PermissionMode, tools []tool.Tool, _ []string) (SessionEngineResult, error) { + return brokerEngineResultWithStoreAndTools(store, tools), nil }, }) if err != nil { @@ -383,12 +557,12 @@ func TestIdleSessionBrokerRefresh_Scenario1_ExistingConversationControls(t *test // AC1.2/AC1.3: starting a refresh removes the active broker wrappers and // invalidates the completed catalogue before the replacement begins, and a // failure at any of the named boundaries (broker reset, engine build, -// attachment commit, snapshot persistence) — or a terminal denial — leaves no +// attachment commit, durable persistence) — or a terminal denial — leaves no // stale or partial broker wrapper executable, without restoring the old // bundle, retaining a partial replacement, or auto-retrying. Each subtest // injects the fault at a different boundary and then proves a subsequent // explicit retry succeeds cleanly. -func TestIdleSessionBrokerRefresh_Scenario1_DestructiveCatalogueReplacement(t *testing.T) { +func TestWorkspaceEnrollmentAuthority_Scenario2_FailureIsAtomic(t *testing.T) { t.Run("BrokerResetFailureLeavesCurrentEnrollmentUndisturbed", func(t *testing.T) { // The only failure branch inside the real Attachment.ResetWorkspaceEnrollment // (engine/adapter/mcpbroker) returns BEFORE any mutation, so a reset @@ -499,7 +673,7 @@ func TestIdleSessionBrokerRefresh_Scenario1_DestructiveCatalogueReplacement(t *t } }) - t.Run("SnapshotPersistenceFailureLeavesNoPartialAuthority", func(t *testing.T) { + t.Run("DurablePersistenceFailureLeavesNoPartialAuthority", func(t *testing.T) { inner := memstore.New() failing := &failNextAuthorizationSaveStore{SessionStore: inner} f := newDestructiveReplacementFixtureWithStore(t, failing) @@ -509,6 +683,13 @@ func TestIdleSessionBrokerRefresh_Scenario1_DestructiveCatalogueReplacement(t *t if err != nil { t.Fatal(err) } + beforeSession, err := inner.Load(t.Context(), f.created.ID) + if err != nil { + t.Fatal(err) + } + beforeAuthority, beforeBound := beforeSession.BoundAuthority() + beforeKeys, beforeKeysPresent := beforeSession.WorkspaceEnrollmentBrokerKeys() + beforePending, beforePendingPresent := beforeSession.PendingWorkspaceEnrollment() t1, err := brokercontract.NewWorkspaceCatalogue(started.Ref, []tool.Tool{enrollmentTool{name: "mcp__slack__post"}}) if err != nil { t.Fatal(err) @@ -516,7 +697,27 @@ func TestIdleSessionBrokerRefresh_Scenario1_DestructiveCatalogueReplacement(t *t f.broker.attachment.result = brokercontract.WorkspaceEnrollmentResult{Ref: started.Ref, Status: brokercontract.WorkspaceEnrollmentConnected, Catalogue: t1} failing.failures.Store(1) if _, err := f.svc.ConnectWorkspaceServices(t.Context(), f.created.ID); !errors.Is(err, ErrInternal) { - t.Fatalf("completion despite a failed snapshot save = %v, want ErrInternal", err) + t.Fatalf("completion despite a failed durable save = %v, want ErrInternal", err) + } + afterFailure, err := inner.Load(t.Context(), f.created.ID) + if err != nil { + t.Fatal(err) + } + afterAuthority, afterBound := afterFailure.BoundAuthority() + afterKeys, afterKeysPresent := afterFailure.WorkspaceEnrollmentBrokerKeys() + afterPending, afterPendingPresent := afterFailure.PendingWorkspaceEnrollment() + if afterBound != beforeBound || !reflect.DeepEqual(afterAuthority, beforeAuthority) || + afterKeysPresent != beforeKeysPresent || !slices.Equal(afterKeys, beforeKeys) || + afterPendingPresent != beforePendingPresent || afterPending != beforePending { + t.Fatalf("AC2.5 durable state after failed completion = authority:%#v bound:%t keys:%q present:%t pending:%#v pendingPresent:%t; want authority:%#v bound:%t keys:%q present:%t pending:%#v pendingPresent:%t", + afterAuthority, afterBound, afterKeys, afterKeysPresent, afterPending, afterPendingPresent, + beforeAuthority, beforeBound, beforeKeys, beforeKeysPresent, beforePending, beforePendingPresent) + } + f.svc.mu.Lock() + candidatePublished := f.svc.sessionEngines[f.created.ID] != nil + f.svc.mu.Unlock() + if candidatePublished { + t.Fatal("AC2.5 save failure left the candidate runtime published") } if pending, ok := f.pendingID(t); !ok || pending != started.Ref.ID { t.Fatalf("pending after a save failure = (%v, %v), want unchanged", pending, ok) @@ -604,7 +805,7 @@ func TestInvariant_idle_session_broker_refresh_serialization(t *testing.T) { // live broker tools and without mecatl ever attaching on its own; the owner's // explicit refresh is the only path that creates a fresh attachment and // installs the exact new admitted names. -func TestIdleSessionBrokerRefresh_Scenario2_RestartAndExplicitRecovery(t *testing.T) { +func TestWorkspaceEnrollmentAuthority_Scenario2_RestartDoesNotResurrectBrokerRuntime(t *testing.T) { store := memstore.New() sessionID := session.SessionID("restart-recovery-session") @@ -625,8 +826,8 @@ func TestIdleSessionBrokerRefresh_Scenario2_RestartAndExplicitRecovery(t *testin SessionEngine: func(context.Context, ProviderSelector, []mcp.ServerConfig, SessionProfile, string, session.PermissionMode) (SessionEngineResult, error) { return brokerEngineResultWithStore(store), nil }, - SessionEngineWithTools: func(context.Context, ProviderSelector, []mcp.ServerConfig, SessionProfile, string, session.PermissionMode, []tool.Tool) (SessionEngineResult, error) { - return brokerEngineResultWithStore(store), nil + SessionEngineWithTools: func(_ context.Context, _ ProviderSelector, _ []mcp.ServerConfig, _ SessionProfile, _ string, _ session.PermissionMode, tools []tool.Tool, _ []string) (SessionEngineResult, error) { + return brokerEngineResultWithStoreAndTools(store, tools), nil }, }) if err != nil { @@ -684,8 +885,8 @@ func TestIdleSessionBrokerRefresh_Scenario2_RestartAndExplicitRecovery(t *testin SessionEngine: func(context.Context, ProviderSelector, []mcp.ServerConfig, SessionProfile, string, session.PermissionMode) (SessionEngineResult, error) { return brokerEngineResultWithStore(store), nil }, - SessionEngineWithTools: func(context.Context, ProviderSelector, []mcp.ServerConfig, SessionProfile, string, session.PermissionMode, []tool.Tool) (SessionEngineResult, error) { - return brokerEngineResultWithStore(store), nil + SessionEngineWithTools: func(_ context.Context, _ ProviderSelector, _ []mcp.ServerConfig, _ SessionProfile, _ string, _ session.PermissionMode, tools []tool.Tool, _ []string) (SessionEngineResult, error) { + return brokerEngineResultWithStoreAndTools(store, tools), nil }, }) if err != nil { @@ -775,8 +976,8 @@ func TestInvariant_idle_session_broker_refresh_preserves_toolhive_custody(t *tes RootAuthority: func(session.SessionKind) session.Authority { return session.Authority{CapabilitySet: governance.CapabilitySet{Tools: []string{"mcp__calendar__list"}}, Provenance: "test"} }, - SessionEngineWithTools: func(context.Context, ProviderSelector, []mcp.ServerConfig, SessionProfile, string, session.PermissionMode, []tool.Tool) (SessionEngineResult, error) { - return brokerEngineResult(), nil + SessionEngineWithTools: func(_ context.Context, _ ProviderSelector, _ []mcp.ServerConfig, _ SessionProfile, _ string, _ session.PermissionMode, tools []tool.Tool, _ []string) (SessionEngineResult, error) { + return brokerEngineResultWithTools(tools), nil }, }) if err != nil { @@ -808,8 +1009,8 @@ func TestInvariant_idle_session_broker_refresh_preserves_toolhive_custody(t *tes RootAuthority: func(session.SessionKind) session.Authority { return session.Authority{CapabilitySet: governance.CapabilitySet{Tools: []string{"mcp__calendar__list"}}, Provenance: "test"} }, - SessionEngineWithTools: func(context.Context, ProviderSelector, []mcp.ServerConfig, SessionProfile, string, session.PermissionMode, []tool.Tool) (SessionEngineResult, error) { - return brokerEngineResult(), nil + SessionEngineWithTools: func(_ context.Context, _ ProviderSelector, _ []mcp.ServerConfig, _ SessionProfile, _ string, _ session.PermissionMode, tools []tool.Tool, _ []string) (SessionEngineResult, error) { + return brokerEngineResultWithTools(tools), nil }, }) if err != nil { @@ -1178,8 +1379,8 @@ func TestWorkspaceEnrollmentAcquiresLeaseBeforePersisting(t *testing.T) { RootAuthority: func(session.SessionKind) session.Authority { return session.Authority{CapabilitySet: governance.CapabilitySet{Tools: []string{"mcp__calendar__list"}}, Provenance: "test"} }, - SessionEngineWithTools: func(context.Context, ProviderSelector, []mcp.ServerConfig, SessionProfile, string, session.PermissionMode, []tool.Tool) (SessionEngineResult, error) { - return brokerEngineResult(), nil + SessionEngineWithTools: func(_ context.Context, _ ProviderSelector, _ []mcp.ServerConfig, _ SessionProfile, _ string, _ session.PermissionMode, tools []tool.Tool, _ []string) (SessionEngineResult, error) { + return brokerEngineResultWithTools(tools), nil }, }) if err != nil { @@ -1216,6 +1417,177 @@ func TestWorkspaceEnrollmentAcquiresLeaseBeforePersisting(t *testing.T) { // structural guarantee, which is what actually pins the AC1.4 requirement // ("Retry keeps the session's run-entry/lease serialization continuously ... // so a prompt cannot enter between them"). +func TestWorkspaceEnrollmentAuthority_Scenario1_PreservesSupportedRestrictions(t *testing.T) { + const brokerTool = "mcp__calendar__list" + var executed atomic.Int32 + store := memstore.New() + svc, broker, created := newWorkspaceEnrollmentServiceWithStore(t, false, store) + svc.cfg.SessionEngineWithTools = func(_ context.Context, _ ProviderSelector, _ []mcp.ServerConfig, _ SessionProfile, _ string, _ session.PermissionMode, tools []tool.Tool, _ []string) (SessionEngineResult, error) { + catalog := tool.NewCatalog() + for _, candidate := range tools { + if err := catalog.Register(candidate); err != nil { + return SessionEngineResult{}, err + } + } + return SessionEngineResult{ + Engine: agent.NewEngine(agent.Deps{ + LLM: mockllm.New(mockllm.ToolCallTurn(session.ToolCall{ID: "broker-call", Name: brokerTool, Args: json.RawMessage(`{}`)}), mockllm.TextTurn("done")), + Catalog: catalog, + Policy: permpolicy.NewPolicy([]governance.Rule{{Scope: governance.ScopeUser, Tool: brokerTool, Effect: governance.Deny}}, nil), + Store: store, + }), + BrokerRegistrationKeys: toolNames(tools), + NonBrokerRegistrationKeys: []string{"Read"}, + Close: func() error { return nil }, + }, nil + } + started, err := svc.ConnectWorkspaceServices(t.Context(), created.ID) + if err != nil { + t.Fatal(err) + } + catalogue, err := brokercontract.NewWorkspaceCatalogue(started.Ref, []tool.Tool{enrollmentTool{name: brokerTool, executed: &executed}}) + if err != nil { + t.Fatal(err) + } + broker.attachment.result = brokercontract.WorkspaceEnrollmentResult{Ref: started.Ref, Status: brokercontract.WorkspaceEnrollmentConnected, Catalogue: catalogue} + if _, err := svc.ConnectWorkspaceServices(t.Context(), created.ID); err != nil { + t.Fatalf("complete enrollment: %v", err) + } + run, err := svc.StartRun(t.Context(), created.ID, "use the calendar") + if err != nil { + t.Fatalf("StartRun: %v", err) + } + var denied bool + for event := range run.Events() { + if event.Type == session.EvToolResult && event.ToolResult != nil && event.ToolResult.IsError { + denied = true + } + } + if got := executed.Load(); got != 0 || !denied { + t.Fatalf("configured deny after enrollment: executed=%d denied=%t", got, denied) + } +} + +func TestWorkspaceEnrollmentAuthority_Scenario3_ConflictPresentationAndOwnership(t *testing.T) { + configureCollision := func(svc *Service) { + svc.cfg.SessionEngineWithTools = func(context.Context, ProviderSelector, []mcp.ServerConfig, SessionProfile, string, session.PermissionMode, []tool.Tool, []string) (SessionEngineResult, error) { + return SessionEngineResult{ + Engine: brokerEngineResult().Engine, + BrokerRegistrationKeys: []string{"mcp__prior__tool"}, + NonBrokerRegistrationKeys: []string{"mcp__prior__tool"}, + Close: func() error { return nil }, + }, nil + } + } + t.Run("owner receives only the safe collision", func(t *testing.T) { + svc, broker, created := newWorkspaceEnrollmentHTTPService(t, true) + replacePersistedBrokerLedger(t, svc, created, []string{"mcp__prior__tool"}, true) + configureCollision(svc) + started, err := svc.ConnectWorkspaceServices(session.WithPrincipal(t.Context(), created.Owner), created.ID) + if err != nil { + t.Fatalf("begin enrollment: %v", err) + } + catalogue, err := brokercontract.NewWorkspaceCatalogue(started.Ref, []tool.Tool{enrollmentTool{name: "mcp__prior__tool"}}) + if err != nil { + t.Fatal(err) + } + broker.attachment.result = brokercontract.WorkspaceEnrollmentResult{Ref: started.Ref, Status: brokercontract.WorkspaceEnrollmentConnected, Catalogue: catalogue} + + _, err = svc.ConnectWorkspaceServices(session.WithPrincipal(t.Context(), created.Owner), created.ID) + if !errors.Is(err, ErrFailedPrecondition) { + t.Fatalf("owner collision = %v, want failed precondition", err) + } + for _, want := range []string{"mcp__prior__tool", "correct the broker configuration", "retry"} { + if !strings.Contains(err.Error(), want) { + t.Fatalf("owner collision = %q, want %q", err, want) + } + } + for _, forbidden := range []string{"https://", "endpoint", "callback", "credential", "token", "oauth"} { + if strings.Contains(strings.ToLower(err.Error()), forbidden) { + t.Fatalf("owner collision leaked %q: %q", forbidden, err) + } + } + }) + t.Run("foreign caller is concealed and cannot mutate", func(t *testing.T) { + svc, _, created := newWorkspaceEnrollmentHTTPService(t, true) + replacePersistedBrokerLedger(t, svc, created, []string{"mcp__prior__tool"}, true) + configureCollision(svc) + before, err := svc.cfg.Store.Load(t.Context(), created.ID) + if err != nil { + t.Fatal(err) + } + foreign := &session.Principal{Issuer: created.Owner.Issuer, Subject: "foreign", GrantType: session.GrantTypeUser} + _, err = svc.ConnectWorkspaceServices(session.WithPrincipal(t.Context(), foreign), created.ID) + if !errors.Is(err, ErrNotFound) || strings.Contains(err.Error(), "mcp__prior__tool") { + t.Fatalf("foreign collision = %v, want concealed not-found failure", err) + } + after, err := svc.cfg.Store.Load(t.Context(), created.ID) + if err != nil { + t.Fatal(err) + } + if !reflect.DeepEqual(after, before) { + t.Fatalf("foreign collision mutated stored session: before=%#v after=%#v", before, after) + } + }) +} + +func TestWorkspaceEnrollmentAuthority_Scenario2_LegacySnapshotIsNotRepaired(t *testing.T) { + t.Run("idle legacy session remains untouched", func(t *testing.T) { + svc, broker, created := newWorkspaceEnrollmentHTTPService(t, false) + replacePersistedBrokerLedger(t, svc, created, nil, false) + before, err := svc.cfg.Store.Load(t.Context(), created.ID) + if err != nil { + t.Fatal(err) + } + + _, err = svc.ConnectWorkspaceServices(t.Context(), created.ID) + if !errors.Is(err, ErrFailedPrecondition) { + t.Fatalf("ConnectWorkspaceServices error = %v, want failed precondition", err) + } + after, err := svc.cfg.Store.Load(t.Context(), created.ID) + if err != nil { + t.Fatal(err) + } + if !reflect.DeepEqual(after, before) { + t.Fatalf("legacy refusal rewrote stored session: before=%#v after=%#v", before, after) + } + if broker.attachment != nil { + begin, _, _ := enrollmentCallCounts(broker.attachment) + if begin != 0 { + t.Fatalf("legacy provenance began %d broker enrollments", begin) + } + } + }) + t.Run("completed legacy session is not reopened or saved", func(t *testing.T) { + svc, _, created := newWorkspaceEnrollmentHTTPService(t, false) + legacy := session.New(created.ID, created.Mode, created.EnvironmentRef, created.Limits, created.CreatedAt) + legacy.ExternalBinding = created.ExternalBinding + authority, _ := created.BoundAuthority() + if err := legacy.BindAuthority(authority); err != nil { + t.Fatal(err) + } + if err := legacy.Complete(); err != nil { + t.Fatal(err) + } + if err := svc.cfg.Store.Save(t.Context(), legacy); err != nil { + t.Fatal(err) + } + if _, err := svc.ConnectWorkspaceServices(t.Context(), created.ID); !errors.Is(err, ErrFailedPrecondition) { + t.Fatalf("ConnectWorkspaceServices error = %v, want failed precondition", err) + } + reloaded, err := svc.cfg.Store.Load(t.Context(), created.ID) + if err != nil { + t.Fatal(err) + } + if reloaded.State != session.StateCompleted { + t.Fatalf("legacy completed session was reopened: state=%q", reloaded.State) + } + if keys, present := reloaded.WorkspaceEnrollmentBrokerKeys(); present || len(keys) != 0 { + t.Fatalf("legacy completed session was repaired: ledger=(%q,%t)", keys, present) + } + }) +} + func TestRetryWorkspaceEnrollmentHoldsRunEntryContinuously(t *testing.T) { svc, broker, created := newWorkspaceEnrollmentHTTPService(t, false) @@ -1329,8 +1701,8 @@ func TestWorkspaceEnrollmentStateLossClearsPendingGate(t *testing.T) { RootAuthority: func(session.SessionKind) session.Authority { return session.Authority{CapabilitySet: governance.CapabilitySet{Tools: []string{"mcp__calendar__list"}}, Provenance: "test"} }, - SessionEngineWithTools: func(context.Context, ProviderSelector, []mcp.ServerConfig, SessionProfile, string, session.PermissionMode, []tool.Tool) (SessionEngineResult, error) { - return brokerEngineResult(), nil + SessionEngineWithTools: func(_ context.Context, _ ProviderSelector, _ []mcp.ServerConfig, _ SessionProfile, _ string, _ session.PermissionMode, tools []tool.Tool, _ []string) (SessionEngineResult, error) { + return brokerEngineResultWithTools(tools), nil }, }) if err != nil { @@ -1390,8 +1762,8 @@ func TestWorkspaceEnrollmentCancelHealsAlreadyGoneBrokerTransaction(t *testing.T RootAuthority: func(session.SessionKind) session.Authority { return session.Authority{CapabilitySet: governance.CapabilitySet{Tools: []string{"mcp__calendar__list"}}, Provenance: "test"} }, - SessionEngineWithTools: func(context.Context, ProviderSelector, []mcp.ServerConfig, SessionProfile, string, session.PermissionMode, []tool.Tool) (SessionEngineResult, error) { - return brokerEngineResult(), nil + SessionEngineWithTools: func(_ context.Context, _ ProviderSelector, _ []mcp.ServerConfig, _ SessionProfile, _ string, _ session.PermissionMode, tools []tool.Tool, _ []string) (SessionEngineResult, error) { + return brokerEngineResultWithTools(tools), nil }, }) if err != nil { @@ -1444,13 +1816,16 @@ func TestWorkspaceEnrollmentPublishesFrozenCatalogueBeforePrompt(t *testing.T) { RootAuthority: func(session.SessionKind) session.Authority { return session.Authority{CapabilitySet: governance.CapabilitySet{Tools: []string{"mcp__calendar__list"}}, Provenance: "test"} }, - SessionEngineWithTools: func(_ context.Context, _ ProviderSelector, _ []mcp.ServerConfig, _ SessionProfile, _ string, _ session.PermissionMode, tools []tool.Tool) (SessionEngineResult, error) { + SessionEngineWithTools: func(_ context.Context, _ ProviderSelector, _ []mcp.ServerConfig, _ SessionProfile, _ string, _ session.PermissionMode, tools []tool.Tool, _ []string) (SessionEngineResult, error) { names := make([]string, len(tools)) for i, candidate := range tools { names[i] = candidate.Spec().Name } catalogues = append(catalogues, names) - return brokerEngineResult(), nil + result := brokerEngineResult() + result.BrokerRegistrationKeys = names + result.NonBrokerRegistrationKeys = []string{"Read"} + return result, nil }, }) if err != nil { @@ -1534,3 +1909,187 @@ func TestWorkspaceEnrollmentPublishesFrozenCatalogueBeforePrompt(t *testing.T) { t.Fatalf("authority = %#v, %v", authority, ok) } } + +func TestWorkspaceEnrollmentAuthority_Scenario2_RestartCollisionFailsPrecondition(t *testing.T) { + svc, _, created := newWorkspaceEnrollmentHTTPService(t, false) + replacePersistedBrokerLedger(t, svc, created, []string{"mcp__prior__tool"}, true) + svc.withdrawBrokerEngine(created.ID) + + var closed bool + svc.cfg.SessionEngine = func(context.Context, ProviderSelector, []mcp.ServerConfig, SessionProfile, string, session.PermissionMode) (SessionEngineResult, error) { + return SessionEngineResult{}, nil + } + svc.cfg.SessionEngineWithTools = func(context.Context, ProviderSelector, []mcp.ServerConfig, SessionProfile, string, session.PermissionMode, []tool.Tool, []string) (SessionEngineResult, error) { + return SessionEngineResult{ + Engine: brokerEngineResult().Engine, + NonBrokerRegistrationKeys: []string{"mcp__prior__tool"}, + Close: func() error { closed = true; return nil }, + }, nil + } + loaded, err := svc.cfg.Store.Load(t.Context(), created.ID) + if err != nil { + t.Fatal(err) + } + _, err = svc.rehydrateSession(t.Context(), loaded) + var collision *WorkspaceEnrollmentCollisionError + if !errors.As(err, &collision) || collision.Key != "mcp__prior__tool" || !errors.Is(err, ErrFailedPrecondition) { + t.Fatalf("rehydration collision = %v, want safe failed-precondition prior registration key", err) + } + if !closed { + t.Fatal("rehydration left the unpublished conflicting candidate open") + } + svc.mu.Lock() + published := svc.sessionEngines[created.ID] + svc.mu.Unlock() + if published != nil { + t.Fatal("rehydration published a non-broker implementation under a prior broker key") + } + + t.Run("excluded ledger key does not block rehydration", func(t *testing.T) { + excluded := session.New(created.ID, created.Mode, created.EnvironmentRef, created.Limits, created.CreatedAt) + excluded.ExternalBinding = created.ExternalBinding + if err := excluded.BindAuthority(session.Authority{Provenance: "test"}); err != nil { + t.Fatal(err) + } + if err := excluded.RestoreWorkspaceEnrollmentBrokerKeys([]string{"mcp__prior__tool"}, true); err != nil { + t.Fatal(err) + } + if err := svc.cfg.Store.Save(t.Context(), excluded); err != nil { + t.Fatal(err) + } + svc.withdrawBrokerEngine(created.ID) + if _, err := svc.rehydrateSession(t.Context(), excluded); err != nil { + t.Fatalf("rehydration with excluded ledger key = %v, want success", err) + } + }) +} + +func TestWorkspaceEnrollmentAuthority_Scenario3_ConflictFailsWholeEnrollment(t *testing.T) { + svc, broker, created := newWorkspaceEnrollmentHTTPService(t, false) + replacePersistedBrokerLedger(t, svc, created, []string{"mcp__prior__tool"}, true) + before, err := svc.cfg.Store.Load(t.Context(), created.ID) + if err != nil { + t.Fatal(err) + } + beforeAuthority, _ := before.BoundAuthority() + beforeKeys, beforePresent := before.WorkspaceEnrollmentBrokerKeys() + + svc.cfg.SessionEngineWithTools = func(context.Context, ProviderSelector, []mcp.ServerConfig, SessionProfile, string, session.PermissionMode, []tool.Tool, []string) (SessionEngineResult, error) { + return SessionEngineResult{ + Engine: brokerEngineResult().Engine, + BrokerRegistrationKeys: []string{"mcp__replacement__tool"}, + NonBrokerRegistrationKeys: []string{"mcp__prior__tool"}, + Close: func() error { return nil }, + }, nil + } + started, err := svc.ConnectWorkspaceServices(t.Context(), created.ID) + if err != nil || started.Status != brokercontract.WorkspaceEnrollmentPending { + t.Fatalf("begin enrollment = (%+v, %v)", started, err) + } + catalogue, err := brokercontract.NewWorkspaceCatalogue(started.Ref, []tool.Tool{enrollmentTool{name: "mcp__replacement__tool"}}) + if err != nil { + t.Fatal(err) + } + broker.attachment.result = brokercontract.WorkspaceEnrollmentResult{Ref: started.Ref, Status: brokercontract.WorkspaceEnrollmentConnected, Catalogue: catalogue} + + _, err = svc.ConnectWorkspaceServices(t.Context(), created.ID) + var collision *WorkspaceEnrollmentCollisionError + if !errors.As(err, &collision) || collision.Key != "mcp__prior__tool" || !errors.Is(err, ErrFailedPrecondition) { + t.Fatalf("connected collision = %v, want presentation-safe failed-precondition collision", err) + } + stored, err := svc.cfg.Store.Load(t.Context(), created.ID) + if err != nil { + t.Fatal(err) + } + if _, pending := stored.PendingWorkspaceEnrollment(); pending { + t.Fatal("collision left enrollment pending in durable session") + } + afterAuthority, _ := stored.BoundAuthority() + afterKeys, afterPresent := stored.WorkspaceEnrollmentBrokerKeys() + if !reflect.DeepEqual(afterAuthority, beforeAuthority) || afterPresent != beforePresent || !slices.Equal(afterKeys, beforeKeys) { + t.Fatalf("collision changed durable authority or provenance: authority=%#v keys=%q present=%t", afterAuthority, afterKeys, afterPresent) + } +} + +func TestWorkspaceEnrollmentAuthority_Scenario1_CompletionLedgerOverflowClearsPending(t *testing.T) { + svc, broker, created := newWorkspaceEnrollmentHTTPService(t, false) + original, err := svc.cfg.Store.Load(t.Context(), created.ID) + if err != nil { + t.Fatal(err) + } + excluded := make([]string, 256) + for i := range excluded { + excluded[i] = fmt.Sprintf("broker-excluded-%03d", i) + } + persisted := session.New(created.ID, created.Mode, created.EnvironmentRef, created.Limits, created.CreatedAt) + persisted.ExternalBinding = original.ExternalBinding + if err := persisted.BindAuthority(session.Authority{CapabilitySet: governance.CapabilitySet{Tools: []string{"Read"}}, Provenance: "test"}); err != nil { + t.Fatal(err) + } + if err := persisted.RestoreWorkspaceEnrollmentBrokerKeys(excluded, true); err != nil { + t.Fatal(err) + } + if err := svc.cfg.Store.Save(t.Context(), persisted); err != nil { + t.Fatal(err) + } + + started, err := svc.ConnectWorkspaceServices(t.Context(), created.ID) + if err != nil || started.Status != brokercontract.WorkspaceEnrollmentPending { + t.Fatalf("begin enrollment = (%+v, %v)", started, err) + } + catalogue, err := brokercontract.NewWorkspaceCatalogue(started.Ref, []tool.Tool{enrollmentTool{name: "broker-new"}}) + if err != nil { + t.Fatal(err) + } + broker.attachment.result = brokercontract.WorkspaceEnrollmentResult{Ref: started.Ref, Status: brokercontract.WorkspaceEnrollmentConnected, Catalogue: catalogue} + if _, err := svc.ConnectWorkspaceServices(t.Context(), created.ID); !errors.Is(err, ErrFailedPrecondition) { + t.Fatalf("completion overflow = %v, want ErrFailedPrecondition", err) + } + stored, err := svc.cfg.Store.Load(t.Context(), created.ID) + if err != nil { + t.Fatal(err) + } + if _, pending := stored.PendingWorkspaceEnrollment(); pending { + t.Fatal("completion overflow left enrollment pending in durable session") + } + keys, present := stored.WorkspaceEnrollmentBrokerKeys() + if !present || !slices.Equal(keys, excluded) { + t.Fatalf("completion overflow changed durable provenance: keys=%q present=%t", keys, present) + } +} + +func TestWorkspaceEnrollmentAuthority_Scenario1_DirectRefreshRefusesBrokerSession(t *testing.T) { + svc, _, created := newWorkspaceEnrollmentHTTPService(t, false) + if created.ExternalBinding == "" { + t.Fatal("fixture session is not broker-bound") + } + replacePersistedBrokerLedger(t, svc, created, []string{"mcp__prior__tool"}, true) + before, err := svc.cfg.Store.Load(t.Context(), created.ID) + if err != nil { + t.Fatal(err) + } + beforeAuthority, _ := before.BoundAuthority() + beforeKeys, beforePresent := before.WorkspaceEnrollmentBrokerKeys() + + var reconciled int + svc.cfg.MCPRefresh = func(context.Context) (MCPRefreshSnapshot, error) { + reconciled++ + return MCPRefreshSnapshot{Revision: 1, Changed: true, ToolNames: []string{"mcp__direct__tool"}}, nil + } + _, err = svc.RefreshMcpSources(t.Context(), created.ID) + if !errors.Is(err, ErrFailedPrecondition) || !strings.Contains(err.Error(), "broker-bound sessions use workspace enrollment") { + t.Fatalf("direct refresh on broker session = %v, want broker-bound failed precondition", err) + } + if reconciled != 0 { + t.Fatalf("direct refresh reconciled %d times for a broker-bound session", reconciled) + } + after, err := svc.cfg.Store.Load(t.Context(), created.ID) + if err != nil { + t.Fatal(err) + } + afterAuthority, _ := after.BoundAuthority() + afterKeys, afterPresent := after.WorkspaceEnrollmentBrokerKeys() + if !reflect.DeepEqual(afterAuthority, beforeAuthority) || afterPresent != beforePresent || !slices.Equal(afterKeys, beforeKeys) { + t.Fatalf("direct refresh changed broker session: authority=%#v keys=%q present=%t", afterAuthority, afterKeys, afterPresent) + } +} diff --git a/internal/app/build.go b/internal/app/build.go index 99f4e68a69..117f2948a4 100644 --- a/internal/app/build.go +++ b/internal/app/build.go @@ -1561,6 +1561,15 @@ func validateMCPAuthority(cfg Config) error { return nil } +// clearGlobalMCPForBroker enforces the broker authority boundary before the +// candidate catalog is built: broker tools are supplied explicitly by the session +// attachment, so neither the static MCP servers nor ToolHive may add a global +// fallback (including a second CallMcpWithQuery). +func clearGlobalMCPForBroker(cfg *Config) { + cfg.MCPServers = nil + cfg.ToolHiveEnabled = false +} + // Build assembles the provider registry, catalog, policy, and engine into a // server.Service per the given Config. It is the single composition root every // cmd/ main calls. @@ -2171,8 +2180,7 @@ func Build(ctx context.Context, cfg Config) (*Built, error) { if brokerSelected { // Authority is exclusive: broker sessions receive only their explicit // attachment wrappers, never the process-global MCP manager as a fallback. - cfg.MCPServers = nil - cfg.ToolHiveEnabled = false + clearGlobalMCPForBroker(&cfg) } cfg.guardrailDetails = server.NewReviewDetailRegistry() cfg.guardrailHealth = &guardrailRouteHealth{} @@ -3297,6 +3305,24 @@ func mountedClientMCPNames(mgr *mcp.Manager) []string { } return names } +func sessionClientMCPManager(ctx context.Context, cfg Config, specs []mcp.ServerConfig) *mcp.Manager { + if len(specs) == 0 { + return nil + } + onError := func(sc mcp.ServerConfig, err error) { + cfg.diag().Log(ctx, port.LevelWarn, "client MCP server unreachable; skipping for this session", + "server", sc.Name, "url", mcp.RedactURL(sc.URL), "err", mcp.RedactError(err)) + } + mgr, err := mcp.NewManager(ctx, specs, onError, cfg.diag()) + if err != nil { + // Client MCP is best-effort for ACP: retain a usable core-tools session + // when every requested server fails. Wire callers enforce their own + // all-or-nothing policy from the mounted-server report. + cfg.diag().Log(ctx, port.LevelWarn, "client MCP: no servers connected for this session; mounting core tools only", + "err", mcp.RedactError(err)) + } + return mgr +} func sessionEngineFactory( cfg Config, @@ -3312,7 +3338,7 @@ func sessionEngineFactory( ) server.SessionEngineFactory { withTools := sessionEngineFactoryWithTools(cfg, reg, provider, store, policy, hooks, mcpProvider, instructions, assets, guardrailWaiver) return func(ctx context.Context, sel server.ProviderSelector, specs []mcp.ServerConfig, profile server.SessionProfile, workspace string, mode session.PermissionMode) (server.SessionEngineResult, error) { - return withTools(ctx, sel, specs, profile, workspace, mode, nil) + return withTools(ctx, sel, specs, profile, workspace, mode, nil, nil) } } @@ -3342,7 +3368,7 @@ func sessionEngineFactoryWithTools( assets catalogAssets, guardrailWaiver *modelhook.WaiverHolder, ) server.SessionEngineWithToolsFactory { - build := func(ctx context.Context, sel server.ProviderSelector, specs []mcp.ServerConfig, profile server.SessionProfile, workspace string, mode session.PermissionMode, sessionTools []tool.Tool) (server.SessionEngineResult, error) { + build := func(ctx context.Context, sel server.ProviderSelector, specs []mcp.ServerConfig, profile server.SessionProfile, workspace string, mode session.PermissionMode, sessionTools []tool.Tool, sessionToolKeys []string) (server.SessionEngineResult, error) { cfg := cfg // Permission resolver overrides belong to this invocation, not the shared closure. runtimeAssets := mcpCatalogAssets(ctx, assets) runtimeRevision := mcpRuntimeRevision(ctx) @@ -3454,46 +3480,7 @@ func sessionEngineFactoryWithTools( // resolves to the 128k floor (inside the resolver). Override still wins. windowFn := reg.windowResolver(cfg, resolvedProviderID, resolvedModel) - onError := func(sc mcp.ServerConfig, err error) { - // REDACTED url (CWE-532). The header map is secret-shaped and never logged, - // but a credential can also ride the URL itself — "?access_token=..." — and - // logging sc.URL verbatim put it in the operator's diagnostics. userinfo is - // now rejected outright by ValidateClientURL; a query-string token cannot be - // (it is indistinguishable from an ordinary parameter), so the URL is - // redacted at the log site instead of trusted to be clean. mcp.RedactURL is - // the SAME policy the validator's own messages use — one redactor, so the - // log and the error cannot drift. - // BOTH the url AND the err must be redacted, and the err is the one that - // was missed: net/http embeds the complete request URL in a connection - // error, and the MCP SDK formats it into its own message text, so a - // query-string token reached the operator log through `err` even though - // `url` beside it was clean. Redacting one of two channels is not - // redacting. - cfg.diag().Log(ctx, port.LevelWarn, "client MCP server unreachable; skipping for this session", - "server", sc.Name, "url", mcp.RedactURL(sc.URL), "err", mcp.RedactError(err)) - } - var mgr *mcp.Manager - if len(specs) > 0 { - m, err := mcp.NewManager(ctx, specs, onError, cfg.diag()) - if err != nil { - // Best-effort HERE, by design: every server failed and the session still - // gets a usable engine (core tools only) rather than failing outright. - // This is the ACP contract — an editor's flaky MCP server should not cost - // the user their session, and ACP has its own channel to say so. - // - // It is NOT the wire contract. A gRPC/HTTP CreateSession caller cannot - // see this WARN, so the Service enforces all-or-nothing on that path - // using MountedClientMCP below. Reporting which servers connected is - // this factory's job; deciding whether a partial mount is acceptable - // belongs to the caller, and the two callers disagree. - // Same redaction as onError above: this error is NewManager's lastErr, - // so it is one of the per-server transport errors and carries that - // server's full URL. - cfg.diag().Log(ctx, port.LevelWarn, "client MCP: no servers connected for this session; mounting core tools only", - "err", mcp.RedactError(err)) - } - mgr = m - } + mgr := sessionClientMCPManager(ctx, cfg, specs) mountedClientMCP := mountedClientMCPNames(mgr) // Assemble the per-session catalog through the SAME assembleCatalog the @@ -3512,7 +3499,7 @@ func sessionEngineFactoryWithTools( // Skill tool. A failed authoritative read clears only these partitions. skillPartitions := hydrateLearnedSkillPartitions(ctx, cfg, runtimeAssets, projectWorkspace) - cat, closeFn, clientToolNames := assembleCatalog(ctx, cfg, reg, store, hooks, &runtimeAssets, catalogSession{ + cat, closeFn, clientToolNames, brokerRegistrationKeys, nonBrokerRegistrationKeys, err := assembleCatalogDetailed(ctx, cfg, reg, store, hooks, &runtimeAssets, catalogSession{ provider: resolvedProvider, providerID: resolvedProviderID, model: resolvedModel, @@ -3523,7 +3510,11 @@ func sessionEngineFactoryWithTools( mode: mode, skillPartitions: skillPartitions, sessionTools: sessionTools, + sessionToolKeys: sessionToolKeys, }) + if err != nil { + return server.SessionEngineResult{}, err + } // Identical to the main engine in every NON-provider Deps field except the // catalog (which carries the extra client MCP + per-session sub-agent tools): @@ -3640,9 +3631,11 @@ func sessionEngineFactoryWithTools( // requested). The factory REPORTS; the Service decides whether a partial // mount is acceptable, because its two callers disagree — see the // best-effort comment on the NewManager error above. - MountedClientMCP: mountedClientMCP, - MountedClientMCPTools: clientToolNames, - Close: closeFn, + MountedClientMCP: mountedClientMCP, + MountedClientMCPTools: clientToolNames, + BrokerRegistrationKeys: brokerRegistrationKeys, + NonBrokerRegistrationKeys: nonBrokerRegistrationKeys, + Close: closeFn, }, nil } return pinMCPRuntimeFactory(assets, build) @@ -4529,7 +4522,7 @@ func buildEngine(ctx context.Context, cfg Config, reg *providerRegistry, provide // not per per-session stop (one of the two sanctioned per-session deltas). assets.sessionFactoryWithTools = sessionEngineFactoryWithTools(cfg, reg, provider, engineStore, sharedPolicy, hooks, mcpProvider, instructions, assets, guardrailWaiver) if cfg.harnessResolver != nil { - assets.sessionContextFactory = func(ctx context.Context, id session.SessionID, owner *session.Principal, acquire server.ExecutionWorkspaceAcquirer, selector server.ProviderSelector, specs []mcp.ServerConfig, profile server.SessionProfile, workspace string, mode session.PermissionMode, extra []tool.Tool, forkSource session.SessionID) (server.SessionEngineResult, error) { + assets.sessionContextFactory = func(ctx context.Context, id session.SessionID, owner *session.Principal, acquire server.ExecutionWorkspaceAcquirer, selector server.ProviderSelector, specs []mcp.ServerConfig, profile server.SessionProfile, workspace string, mode session.PermissionMode, extra []tool.Tool, sessionToolKeys []string, forkSource session.SessionID) (server.SessionEngineResult, error) { // Reject an invalid selector before principal-scoped source binding creates // unpublished state that cannot be adopted by a session. if selector.ProviderID != "" { @@ -4566,7 +4559,7 @@ func buildEngine(ctx context.Context, cfg Config, reg *providerRegistry, provide } sessionCfg.skillCommandInputs = skillCommandInputs{metas: seam.metas, source: seam.source} factory := sessionEngineFactoryWithTools(sessionCfg, reg, provider, engineStore, sharedPolicy, hooks, mcpProvider, replaceHarnessInstructions(instructions, sessionCfg), sessionAssets, guardrailWaiver) - result, err := factory(ctx, selector, specs, profile, workspace, mode, extra) + result, err := factory(ctx, selector, specs, profile, workspace, mode, extra, sessionToolKeys) if err != nil { release() return result, err diff --git a/internal/app/catalog.go b/internal/app/catalog.go index a102930ebc..5013aa5da5 100644 --- a/internal/app/catalog.go +++ b/internal/app/catalog.go @@ -23,6 +23,7 @@ package app import ( "context" + "fmt" "strings" agents "github.com/stacklok/mecatl/engine/adapter/agentfs" @@ -192,8 +193,11 @@ type catalogSession struct { // share one principal-scoped catalog selection. skillPartitions []learning.SkillPartition // sessionTools are explicit wrappers owned by one host attachment. They are - // never recovered from context values or a global MCP manager. - sessionTools []tool.Tool + // never recovered from context values or a global MCP manager. sessionToolKeys, + // when supplied, are the broker's already-frozen registration keys paired by + // index with sessionTools. + sessionTools []tool.Tool + sessionToolKeys []string } // assembleCatalog registers every tool family into a fresh catalog, in the @@ -218,7 +222,7 @@ type catalogSession struct { // into the returned catalog (see the clientToolNames comment inline below for // why it must come from the catalog itself rather than a second manager // read) -- nil whenever s.clientMgr is nil, i.e. every build-time/shared call. -func assembleCatalog(ctx context.Context, cfg Config, reg *providerRegistry, store port.SessionStore, hooks port.HookRunner, a *catalogAssets, s catalogSession) (*tool.Catalog, func() error, []string) { +func assembleCatalogDetailed(ctx context.Context, cfg Config, reg *providerRegistry, store port.SessionStore, hooks port.HookRunner, a *catalogAssets, s catalogSession) (*tool.Catalog, func() error, []string, []string, []string, error) { if profile, ok := a.userModelStore.(prompt.OperatorProfileSource); ok { cfg.operatorProfileSource = profile } @@ -227,16 +231,6 @@ func assembleCatalog(ctx context.Context, cfg Config, reg *providerRegistry, sto classified.captureEach(coreToolClassification, func() { registerCoreTools(cfg, cat, s.narrate, s.noFS, s.remote, a.searchProvider) }) - for _, sessionTool := range s.sessionTools { - // A direct global manager retains ownership of its existing query wrapper. - // The attachment-bound variant is for broker-only session catalogues. - if sessionTool.Spec().Name == "CallMcpWithQuery" && a.globalMgr != nil { - continue - } - classified.mustRegister(sessionTool, classification(server.KindDerived, - "session-bound wrapper supplied explicitly by the host attachment")) - } - for _, extra := range cfg.extraCoreTools { entry, ok := cfg.extraCoreToolClassifications[extra.Spec().Name] if !ok { @@ -326,6 +320,22 @@ func assembleCatalog(ctx context.Context, cfg Config, reg *providerRegistry, sto registerSkillFamily(ctx, cfg, cat, *a, s) }) + brokerRegistrationKeys, err := registerSessionTools(classified, s.sessionTools, s.sessionToolKeys, clientClose, subagentClose) + if err != nil { + return nil, nil, nil, nil, nil, err + } + + brokerSet := make(map[string]struct{}, len(brokerRegistrationKeys)) + for _, name := range brokerRegistrationKeys { + brokerSet[name] = struct{}{} + } + nonBrokerRegistrationKeys := make([]string, 0, len(cat.Names())-len(brokerRegistrationKeys)) + for _, name := range cat.Names() { + if _, broker := brokerSet[name]; !broker { + nonBrokerRegistrationKeys = append(nonBrokerRegistrationKeys, name) + } + } + if cfg.catalogClassificationObserver != nil { cfg.catalogClassificationObserver(classified.snapshot()) } @@ -334,6 +344,90 @@ func assembleCatalog(ctx context.Context, cfg Config, reg *providerRegistry, sto if closeFn == nil { closeFn = func() error { return nil } } + return cat, closeFn, clientToolNames, brokerRegistrationKeys, nonBrokerRegistrationKeys, nil +} + +func registerSessionTools(classified *classifiedCatalog, sessionTools []tool.Tool, sessionToolKeys []string, clientClose, subagentClose func() error) ([]string, error) { + cleanup := func() { + if clientClose != nil { + _ = clientClose() + } + if subagentClose != nil { + _ = subagentClose() + } + } + frozenKeys := sessionToolKeys != nil + if frozenKeys { + if len(sessionToolKeys) != len(sessionTools) { + cleanup() + return nil, fmt.Errorf("session tool registration keys do not match tools") + } + for _, key := range sessionToolKeys { + if _, exists := classified.catalog.Lookup(key); exists { + cleanup() + return nil, &server.WorkspaceEnrollmentCollisionError{Key: key} + } + } + } + + brokerKeys := make([]string, 0, len(sessionTools)) + initialNames := classified.catalog.Names() + runningNames := make(map[string]struct{}, len(initialNames)) + for _, name := range initialNames { + runningNames[name] = struct{}{} + } + for i, sessionTool := range sessionTools { + registeredKey := "" + if frozenKeys { + registeredKey = sessionToolKeys[i] + if _, exists := runningNames[registeredKey]; exists { + cleanup() + return nil, &server.WorkspaceEnrollmentCollisionError{Key: registeredKey} + } + } + if err := classified.catalog.Register(sessionTool); err != nil { + cleanup() + return nil, fmt.Errorf("register session tool: %w", err) + } + if frozenKeys { + if _, exists := classified.catalog.Lookup(registeredKey); !exists { + cleanup() + return nil, fmt.Errorf("session tool registered key %q does not match frozen key", registeredKey) + } + runningNames[registeredKey] = struct{}{} + brokerKeys = append(brokerKeys, registeredKey) + classified.entries[registeredKey] = server.ClassificationEntry{Kind: server.KindDerived, + Rationale: "session-bound wrapper supplied explicitly by the host attachment"} + } + } + if frozenKeys { + return brokerKeys, nil + } + + for _, name := range classified.catalog.Names() { + if _, existed := runningNames[name]; existed { + continue + } + runningNames[name] = struct{}{} + brokerKeys = append(brokerKeys, name) + classified.entries[name] = server.ClassificationEntry{Kind: server.KindDerived, + Rationale: "session-bound wrapper supplied explicitly by the host attachment"} + } + if len(brokerKeys) != len(sessionTools) { + cleanup() + return nil, fmt.Errorf("session tool registered unexpected key count") + } + return brokerKeys, nil +} + +// assembleCatalog is the static-registration convenience used by build-time and +// test callers. Dynamic host tool registration uses assembleCatalogDetailed so +// duplicate registration is an ordinary candidate-build failure, never a panic. +func assembleCatalog(ctx context.Context, cfg Config, reg *providerRegistry, store port.SessionStore, hooks port.HookRunner, a *catalogAssets, s catalogSession) (*tool.Catalog, func() error, []string) { + cat, closeFn, clientToolNames, _, _, err := assembleCatalogDetailed(ctx, cfg, reg, store, hooks, a, s) + if err != nil { + panic(err) + } return cat, closeFn, clientToolNames } diff --git a/internal/app/catalog_classification.go b/internal/app/catalog_classification.go index 6488378d98..02861559dd 100644 --- a/internal/app/catalog_classification.go +++ b/internal/app/catalog_classification.go @@ -24,8 +24,7 @@ func newClassifiedCatalog() *classifiedCatalog { } func (c *classifiedCatalog) classifyAdded(before map[string]struct{}, entry server.ClassificationEntry) { - for _, registered := range c.catalog.Tools() { - name := registered.Spec().Name + for _, name := range c.catalog.Names() { if _, existed := before[name]; !existed { c.entries[name] = entry } @@ -41,11 +40,11 @@ func (c *classifiedCatalog) capture(entry server.ClassificationEntry, register f func (c *classifiedCatalog) captureEach(classify func(tool.Tool) (server.ClassificationEntry, bool), register func()) { before := c.names() register() - for _, registered := range c.catalog.Tools() { - name := registered.Spec().Name + for _, name := range c.catalog.Names() { if _, existed := before[name]; existed { continue } + registered, _ := c.catalog.Lookup(name) if entry, ok := classify(registered); ok { c.entries[name] = entry } @@ -69,9 +68,9 @@ func (c *classifiedCatalog) mustRegister(t tool.Tool, entry *server.Classificati } func (c *classifiedCatalog) names() map[string]struct{} { - names := make(map[string]struct{}, len(c.catalog.Tools())) - for _, registered := range c.catalog.Tools() { - names[registered.Spec().Name] = struct{}{} + names := make(map[string]struct{}, len(c.catalog.Names())) + for _, name := range c.catalog.Names() { + names[name] = struct{}{} } return names } @@ -85,12 +84,7 @@ func (c *classifiedCatalog) snapshot() map[string]server.ClassificationEntry { } func (c *classifiedCatalog) validate(surface string) error { - tools := c.catalog.Tools() - names := make([]string, 0, len(tools)) - for _, registered := range tools { - names = append(names, registered.Spec().Name) - } - if errs := server.ValidateClassifiedNames(surface, c.entries, names); len(errs) > 0 { + if errs := server.ValidateClassifiedNames(surface, c.entries, c.catalog.Names()); len(errs) > 0 { return fmt.Errorf("caller-separation tool classification: %w", errors.Join(errs...)) } return nil diff --git a/internal/app/catalog_classification_test.go b/internal/app/catalog_classification_test.go index fdf48ee1eb..8b5e2e43e1 100644 --- a/internal/app/catalog_classification_test.go +++ b/internal/app/catalog_classification_test.go @@ -2,7 +2,9 @@ package app import ( "context" + "errors" "fmt" + "reflect" "strings" "testing" @@ -11,6 +13,7 @@ import ( "github.com/stacklok/mecatl/engine/learning" "github.com/stacklok/mecatl/engine/session" "github.com/stacklok/mecatl/engine/tool" + "github.com/stacklok/mecatl/internal/adapter/mcpauthority" "github.com/stacklok/mecatl/internal/adapter/server" ) @@ -22,6 +25,90 @@ func (classificationNamedTool) Execute(_ context.Context, call session.ToolCall, return session.NewToolResult(call.ID, "ok"), nil } +type countingRegistrationTool struct { + name string + calls int +} + +func (t *countingRegistrationTool) Spec() tool.ToolSpec { + t.calls++ + return tool.ToolSpec{Name: t.name} +} +func (*countingRegistrationTool) ReadOnly() bool { return true } +func (*countingRegistrationTool) Execute(context.Context, session.ToolCall, tool.Environment) (session.ToolResult, error) { + return session.ToolResult{}, nil +} + +func TestBrokerSelectionBuildExcludesGlobalMCPInputs(t *testing.T) { + built, err := buildIsolated(t, t.Context(), Config{ + Workspace: t.TempDir(), + UseMock: true, + NoSoul: true, + MCPAuthority: mcpauthority.NewBroker(mcpauthority.BrokerConfig{}), + ToolHiveEnabled: true, + }) + if err != nil { + t.Fatalf("Build: %v", err) + } + defer built.Close() + + // ToolHive source discovery is only selected when ToolHiveEnabled reaches the + // real Build path. Broker mode must expose neither that global source nor its + // global-tool fallback; broker tools arrive only through an attachment. + for _, source := range built.Service.ListMcpSources(t.Context()).Sources { + if source.Kind == "toolhive" { + t.Fatalf("broker Build retained the global ToolHive source: %+v", source) + } + } +} + +func TestWorkspaceEnrollmentAuthority_Scenario3_RegistrationMetadataAndDuplicateHandoff(t *testing.T) { + t.Run("exact broker and non-broker registration handoff", func(t *testing.T) { + classified := newClassifiedCatalog() + classified.catalog.MustRegister(classificationNamedTool("Read")) + broker := &countingRegistrationTool{name: "mcp__calendar__list"} + keys, err := registerSessionTools(classified, []tool.Tool{broker}, []string{"mcp__calendar__list"}, nil, nil) + if err != nil { + t.Fatalf("registerSessionTools: %v", err) + } + if broker.calls != 1 || len(keys) != 1 || keys[0] != "mcp__calendar__list" { + t.Fatalf("broker registration = keys %q, Spec calls %d", keys, broker.calls) + } + if got := classified.catalog.Names(); len(got) != 2 || got[0] != "Read" || got[1] != "mcp__calendar__list" { + t.Fatalf("final catalog names = %q", got) + } + }) + t.Run("incoming duplicate is safe and does not retry Spec", func(t *testing.T) { + classified := newClassifiedCatalog() + classified.catalog.MustRegister(classificationNamedTool("Read")) + broker := &countingRegistrationTool{name: "Read"} + _, err := registerSessionTools(classified, []tool.Tool{broker}, []string{"Read"}, nil, nil) + var collision *server.WorkspaceEnrollmentCollisionError + if !errors.As(err, &collision) || collision.Key != "Read" { + t.Fatalf("duplicate handoff = %v, want safe Read collision", err) + } + if broker.calls != 0 { + t.Fatalf("preflight called Spec %d times, want none", broker.calls) + } + }) + t.Run("unfrozen registration captures each key without an extra Spec call", func(t *testing.T) { + classified := newClassifiedCatalog() + classified.catalog.MustRegister(classificationNamedTool("Read")) + first := &countingRegistrationTool{name: "mcp__calendar__list"} + second := &countingRegistrationTool{name: "mcp__calendar__get"} + keys, err := registerSessionTools(classified, []tool.Tool{first, second}, nil, nil, nil) + if err != nil { + t.Fatalf("registerSessionTools: %v", err) + } + if first.calls != 1 || second.calls != 1 { + t.Fatalf("Spec calls = first %d, second %d; want one each", first.calls, second.calls) + } + if want := []string{"mcp__calendar__get", "mcp__calendar__list"}; !reflect.DeepEqual(keys, want) { + t.Fatalf("registration keys = %q, want %q", keys, want) + } + }) +} + func TestCallerSeparation_ClassificationKindsMatchRealRegistrationContext(t *testing.T) { var got map[string]server.ClassificationEntry built, err := buildIsolated(t, context.Background(), Config{ diff --git a/internal/app/harness_context_execution_source_test.go b/internal/app/harness_context_execution_source_test.go index 73d09d7d7b..eca58aefab 100644 --- a/internal/app/harness_context_execution_source_test.go +++ b/internal/app/harness_context_execution_source_test.go @@ -501,7 +501,7 @@ func testExecutionSourceFailureBoundaries(t *testing.T) { } defer resolver.Close() engine := agent.NewEngine(agent.Deps{LLM: mockllm.New(mockllm.TextTurn("done")), Catalog: tool.NewCatalog()}) - factory := func(factoryCtx context.Context, id session.SessionID, owner *session.Principal, acquire server.ExecutionWorkspaceAcquirer, _ server.ProviderSelector, _ []mcp.ServerConfig, profile server.SessionProfile, _ string, _ session.PermissionMode, _ []tool.Tool, _ session.SessionID) (server.SessionEngineResult, error) { + factory := func(factoryCtx context.Context, id session.SessionID, owner *session.Principal, acquire server.ExecutionWorkspaceAcquirer, _ server.ProviderSelector, _ []mcp.ServerConfig, profile server.SessionProfile, _ string, _ session.PermissionMode, _ []tool.Tool, _ []string, _ session.SessionID) (server.SessionEngineResult, error) { _, release, err := resolver.BorrowWithExecutionWorkspace(factoryCtx, id, owner, string(profile), acquire) if err != nil { return server.SessionEngineResult{}, err diff --git a/internal/app/mcp_broker_p10_test.go b/internal/app/mcp_broker_p10_test.go index b6b93f9b2a..fddf9cc90d 100644 --- a/internal/app/mcp_broker_p10_test.go +++ b/internal/app/mcp_broker_p10_test.go @@ -6,6 +6,7 @@ import ( "encoding/json" "net/http" "net/http/httptest" + "sync/atomic" "testing" "time" @@ -58,7 +59,7 @@ func TestCallMcpWithQueryBrokerSupport_Scenario1_BrokerOnlyFactory(t *testing.T) t.Fatal("broker attachment does not expose query wrapper") } tools = append(tools, queryAttachment.CallMcpWithQueryTool()) - result, err := factory(context.Background(), server.ProviderSelector{}, nil, server.ProfileDefault, "/workspace", session.ModeDefault, tools) + result, err := factory(context.Background(), server.ProviderSelector{}, nil, server.ProfileDefault, "/workspace", session.ModeDefault, tools, nil) if err != nil { t.Fatal(err) } @@ -92,7 +93,7 @@ func TestCallMcpWithQueryBrokerSupport_Scenario1_BrokerOnlyFactory(t *testing.T) if count != 1 { t.Fatalf("model received %d query wrappers", count) } - empty, err := factory(t.Context(), server.ProviderSelector{}, nil, server.ProfileDefault, "/workspace", session.ModeDefault, nil) + empty, err := factory(t.Context(), server.ProviderSelector{}, nil, server.ProfileDefault, "/workspace", session.ModeDefault, nil, nil) if err != nil { t.Fatal(err) } @@ -167,6 +168,54 @@ func TestCallMcpWithQueryBrokerSupport_BrokerOnlyServiceRegistration(t *testing. } } +func TestWorkspaceEnrollmentAuthority_Scenario1_PreservesConfiguredBrokerDeny(t *testing.T) { + const brokerTool = "mcp__calendar__list" + var calls atomic.Int32 + provider := mockllm.New( + mockllm.ToolCallTurn(session.NewToolCall("broker-call", brokerTool, json.RawMessage(`{}`))), + mockllm.TextTurn("done"), + ) + built, err := buildIsolated(t, t.Context(), Config{ + Workspace: t.TempDir(), + MockProvider: provider, + NoSoul: true, + PermissionConfigs: []string{writeOperatorSettingsFile(t, "permissions:\n deny:\n - "+brokerTool+"\n")}, + MCPAuthority: mcpauthority.NewBroker(mcpauthority.BrokerConfig{ + Routes: []permconfig.MCPServerProfile{{Name: "calendar", Auth: permconfig.MCPAuthProfile{Mode: "none"}}}, + }), + MCPBrokerDiscovered: []mcpbroker.ToolDefinition{{ + Backend: "calendar", Name: brokerTool, Schema: json.RawMessage(`{"type":"object"}`), ReadOnly: true, + }}, + MCPBrokerCaller: func(context.Context, mcpbroker.SessionRef, string, session.ToolCall) (session.ToolResult, error) { + calls.Add(1) + return session.NewToolResult("broker-call", "unexpected execution"), nil + }, + }) + if err != nil { + t.Fatalf("Build: %v", err) + } + defer built.Close() + + sess, err := built.Service.CreateSession(t.Context(), session.ModeDefault, session.Limits{}) + if err != nil { + t.Fatalf("CreateSession: %v", err) + } + run, err := built.Service.StartRun(t.Context(), sess.ID, "use the calendar") + if err != nil { + t.Fatalf("StartRun: %v", err) + } + defer built.Service.FinishRun(sess.ID, run) + for event := range run.Events() { + if event.Type == session.EvToolResult && event.ToolResult != nil && event.ToolResult.IsError { + if got := calls.Load(); got != 0 { + t.Fatalf("configured deny executed broker tool %d times", got) + } + return + } + } + t.Fatal("configured permissions.deny did not reject the broker tool") +} + func TestBuiltMountsFixedMCPBrokerHandlerBundle(t *testing.T) { built := &Built{ MCPBrokerHandlers: mcpbroker.HandlerBundle{Callback: http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { diff --git a/internal/app/mcp_runtime.go b/internal/app/mcp_runtime.go index fedc262957..f84eb325bf 100644 --- a/internal/app/mcp_runtime.go +++ b/internal/app/mcp_runtime.go @@ -254,13 +254,13 @@ func mcpCatalogAssets(ctx context.Context, assets catalogAssets) catalogAssets { } func pinMCPRuntimeFactory(assets catalogAssets, build serveradapter.SessionEngineWithToolsFactory) serveradapter.SessionEngineWithToolsFactory { - return func(ctx context.Context, sel serveradapter.ProviderSelector, specs []mcp.ServerConfig, profile serveradapter.SessionProfile, workspace string, mode session.PermissionMode, sessionTools []tool.Tool) (serveradapter.SessionEngineResult, error) { + return func(ctx context.Context, sel serveradapter.ProviderSelector, specs []mcp.ServerConfig, profile serveradapter.SessionProfile, workspace string, mode session.PermissionMode, sessionTools []tool.Tool, sessionToolKeys []string) (serveradapter.SessionEngineResult, error) { pinned, _, _, release, err := pinCatalogAssets(ctx, assets) if err != nil { return serveradapter.SessionEngineResult{}, err } defer release() - return build(pinned, sel, specs, profile, workspace, mode, sessionTools) + return build(pinned, sel, specs, profile, workspace, mode, sessionTools, sessionToolKeys) } } diff --git a/internal/mcpbroker/workspace_enrollment_test.go b/internal/mcpbroker/workspace_enrollment_test.go index 4a70f195ad..b98d342f94 100644 --- a/internal/mcpbroker/workspace_enrollment_test.go +++ b/internal/mcpbroker/workspace_enrollment_test.go @@ -113,7 +113,7 @@ func TestWorkspaceEnrollmentResultShapes(t *testing.T) { } } -func TestWorkspaceCatalogueIsImmutableFrozenSnapshot(t *testing.T) { +func TestWorkspaceCatalogueFreezesRegistrationKeysAndToolSpecs(t *testing.T) { ref := enrollmentRef() schema := json.RawMessage(`{"type":"object"}`) source := &enrollmentTool{ diff --git a/user-docs/operating/mecak8s/identity-and-client-access.md b/user-docs/operating/mecak8s/identity-and-client-access.md index 0c7c6493a7..840c227850 100644 --- a/user-docs/operating/mecak8s/identity-and-client-access.md +++ b/user-docs/operating/mecak8s/identity-and-client-access.md @@ -221,6 +221,14 @@ placeholders with a frozen per-session catalog. Failure exposes no partial catalog. Keep preregistered client secrets in `SecretKeyRef`; broker metadata and profiles remain non-secret ConfigMap data. +:::caution[Snapshots without provenance need a fresh session] + +A session snapshot without workspace broker provenance cannot identify which +broker tools its authority covered. Mecatl cannot infer or repair that provenance. +Start a fresh session before using `/tools-connect`. + +::: + :::caution[Broker mode is single-replica] Broker sessions and OAuth state are process-local. The chart requires From 38f3f2c34822527f91ee06e06c4e9644fa9e3b2b Mon Sep 17 00:00:00 2001 From: Jakub Hrozek Date: Wed, 14 Oct 2026 18:01:46 +0200 Subject: [PATCH 2/3] test(app): drain run events before FinishRun in broker deny test The test returned on the first denied tool result, so the deferred FinishRun read session state while the engine goroutine was still in BeginTurn on the mock's second turn. Drain the stream first, then assert. Co-Authored-By: mecatl --- internal/app/mcp_broker_p10_test.go | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/internal/app/mcp_broker_p10_test.go b/internal/app/mcp_broker_p10_test.go index fddf9cc90d..12d0784228 100644 --- a/internal/app/mcp_broker_p10_test.go +++ b/internal/app/mcp_broker_p10_test.go @@ -205,15 +205,19 @@ func TestWorkspaceEnrollmentAuthority_Scenario1_PreservesConfiguredBrokerDeny(t t.Fatalf("StartRun: %v", err) } defer built.Service.FinishRun(sess.ID, run) + // Drain to the end: FinishRun must not race the engine's follow-up turn. + denied := false for event := range run.Events() { if event.Type == session.EvToolResult && event.ToolResult != nil && event.ToolResult.IsError { - if got := calls.Load(); got != 0 { - t.Fatalf("configured deny executed broker tool %d times", got) - } - return + denied = true } } - t.Fatal("configured permissions.deny did not reject the broker tool") + if !denied { + t.Fatal("configured permissions.deny did not reject the broker tool") + } + if got := calls.Load(); got != 0 { + t.Fatalf("configured deny executed broker tool %d times", got) + } } func TestBuiltMountsFixedMCPBrokerHandlerBundle(t *testing.T) { From ee4d12679a672af4949c83c835b8ee0d52222036 Mon Sep 17 00:00:00 2001 From: Jakub Hrozek Date: Wed, 14 Oct 2026 18:08:54 +0200 Subject: [PATCH 3/3] test(enrollment): cover workspace enrollment authority regressions Adds server and app tests for frozen-key forwarding, registration collisions, carryover, and collision settlement. Co-Authored-By: mecatl --- .../workspace_enrollment_regression_test.go | 245 ++++++++++++++++++ .../app/workspace_enrollment_context_test.go | 74 ++++++ 2 files changed, 319 insertions(+) create mode 100644 internal/adapter/server/workspace_enrollment_regression_test.go create mode 100644 internal/app/workspace_enrollment_context_test.go diff --git a/internal/adapter/server/workspace_enrollment_regression_test.go b/internal/adapter/server/workspace_enrollment_regression_test.go new file mode 100644 index 0000000000..bd7dd6ab3f --- /dev/null +++ b/internal/adapter/server/workspace_enrollment_regression_test.go @@ -0,0 +1,245 @@ +package server + +import ( + "context" + "errors" + "reflect" + "slices" + "testing" + "time" + + "github.com/stacklok/mecatl/engine/adapter/memstore" + "github.com/stacklok/mecatl/engine/adapter/mockllm" + "github.com/stacklok/mecatl/engine/adapter/permpolicy" + "github.com/stacklok/mecatl/engine/agent" + "github.com/stacklok/mecatl/engine/governance" + "github.com/stacklok/mecatl/engine/session" + "github.com/stacklok/mecatl/engine/tool" + "github.com/stacklok/mecatl/internal/adapter/mcp" + brokercontract "github.com/stacklok/mecatl/internal/mcpbroker" +) + +func TestWorkspaceEnrollmentRegistrationOverlapRequiresAuthority(t *testing.T) { + for _, tc := range []struct { + name string + authority []string + registrations []string + wantOverlap bool + }{ + {name: "authorized overlap", authority: []string{"prior"}, registrations: []string{"prior"}, wantOverlap: true}, + {name: "excluded authority", registrations: []string{"prior"}}, + {name: "unrelated registration", authority: []string{"prior"}, registrations: []string{"other"}}, + } { + t.Run(tc.name, func(t *testing.T) { + key, overlap := workspaceEnrollmentRegistrationOverlap([]string{"prior"}, tc.authority, tc.registrations) + if overlap != tc.wantOverlap || overlap && key != "prior" { + t.Fatalf("overlap = (%q, %t), want (%q, %t)", key, overlap, "prior", tc.wantOverlap) + } + }) + } +} + +func TestCallSessionEnginePreservesFrozenKeyNilness(t *testing.T) { + for _, keys := range [][]string{nil, {}} { + wantNil := keys == nil + var gotNil bool + svc := &Service{cfg: Config{SessionContextEngine: func(_ context.Context, _ session.SessionID, _ *session.Principal, _ ExecutionWorkspaceAcquirer, _ ProviderSelector, _ []mcp.ServerConfig, _ SessionProfile, _ string, _ session.PermissionMode, _ []tool.Tool, got []string, _ session.SessionID) (SessionEngineResult, error) { + gotNil = got == nil + return SessionEngineResult{}, nil + }}} + if _, err := svc.callSessionEngine(t.Context(), "s", nil, nil, ProviderSelector{}, nil, ProfileDefault, "", session.ModeDefault, nil, keys); err != nil { + t.Fatal(err) + } + if gotNil != wantNil { + t.Fatalf("forwarded keys nil = %t, want %t", gotNil, wantNil) + } + } +} + +func TestCreateCarryoverRejectsPerSessionRegistrationCollision(t *testing.T) { + svc, _, source := newWorkspaceEnrollmentHTTPService(t, false) + const key = "mcp__prior__tool" + replacePersistedBrokerLedger(t, svc, source, []string{key}, true) + before, err := svc.cfg.Store.Load(t.Context(), source.ID) + if err != nil { + t.Fatal(err) + } + beforeAuthority, _ := before.BoundAuthority() + beforeKeys, beforePresent := before.WorkspaceEnrollmentBrokerKeys() + + var closed bool + svc.cfg.SessionEngineWithTools = func(context.Context, ProviderSelector, []mcp.ServerConfig, SessionProfile, string, session.PermissionMode, []tool.Tool, []string) (SessionEngineResult, error) { + return SessionEngineResult{ + Engine: brokerEngineResult().Engine, + NonBrokerRegistrationKeys: []string{key}, + Close: func() error { closed = true; return nil }, + }, nil + } + svc.cfg.NewID = func() session.SessionID { return "carryover-collision-target" } + _, err = svc.CreateSessionWithProfile(t.Context(), session.ModeDefault, session.Limits{}, ProviderSelector{}, ProfileDefault, WithSourceSession(source.ID)) + var collision *WorkspaceEnrollmentCollisionError + if !errors.As(err, &collision) || collision.Key != key || !errors.Is(err, ErrFailedPrecondition) { + t.Fatalf("carryover create collision = %v, want typed failed-precondition %q", err, key) + } + if !closed { + t.Fatal("rejected per-session engine candidate was not closed") + } + if stored, loadErr := svc.cfg.Store.Load(t.Context(), "carryover-collision-target"); loadErr == nil || stored != nil { + t.Fatalf("rejected carryover session was persisted: session=%v err=%v", stored, loadErr) + } + after, err := svc.cfg.Store.Load(t.Context(), source.ID) + if err != nil { + t.Fatal(err) + } + afterAuthority, _ := after.BoundAuthority() + afterKeys, afterPresent := after.WorkspaceEnrollmentBrokerKeys() + if !reflect.DeepEqual(afterAuthority, beforeAuthority) || beforePresent != afterPresent || !slices.Equal(beforeKeys, afterKeys) { + t.Fatalf("collision changed source authority/provenance: authority=%#v keys=%q present=%t", afterAuthority, afterKeys, afterPresent) + } +} + +func TestCreateCarryoverSharedEngineChecksOnlyAuthorizedBrokerKeys(t *testing.T) { + for _, tc := range []struct { + name string + authority []string + wantCollision bool + }{ + {name: "authorized key collides", authority: []string{"Read", "mcp__prior__tool"}, wantCollision: true}, + {name: "excluded key is preserved", authority: []string{"Read"}}, + } { + t.Run(tc.name, func(t *testing.T) { + store := memstore.New() + catalog := tool.NewCatalog() + if err := catalog.Register(enrollmentTool{name: "mcp__prior__tool"}); err != nil { + t.Fatal(err) + } + shared := agent.NewEngine(agent.Deps{LLM: mockllm.New(mockllm.TextTurn("done")), Catalog: catalog, Policy: permpolicy.NewPolicy(nil, nil)}) + svc, err := NewService(Config{ + Engine: shared, Store: store, PlacementProvider: brokerPlacementProvider{}, PlacementScope: "test", SharedEngineRoot: "/workspace", + NewID: func() session.SessionID { return "shared-carryover-target" }, + RootAuthority: func(session.SessionKind) session.Authority { + return session.Authority{CapabilitySet: governance.CapabilitySet{Tools: []string{"Read"}}, Provenance: "test"} + }, + }) + if err != nil { + t.Fatal(err) + } + t.Cleanup(svc.Close) + source := session.New("shared-carryover-source", session.ModeDefault, session.EnvironmentRef{Kind: session.EnvKindLocal, ID: "/workspace", Revision: "in-tree-v1"}, session.Limits{}, time.Unix(0, 0)) + if err := source.BindAuthority(session.Authority{CapabilitySet: governance.CapabilitySet{Tools: tc.authority}, Provenance: "test"}); err != nil { + t.Fatal(err) + } + if err := source.RestoreWorkspaceEnrollmentBrokerKeys([]string{"mcp__prior__tool"}, true); err != nil { + t.Fatal(err) + } + if err := store.Create(t.Context(), source); err != nil { + t.Fatal(err) + } + + created, err := svc.CreateSessionWithProfile(t.Context(), session.ModeDefault, session.Limits{}, ProviderSelector{}, ProfileDefault, WithSourceSession(source.ID)) + var collision *WorkspaceEnrollmentCollisionError + if tc.wantCollision { + if !errors.As(err, &collision) || collision.Key != "mcp__prior__tool" || !errors.Is(err, ErrFailedPrecondition) { + t.Fatalf("shared-engine collision = %v, want typed failed-precondition", err) + } + if stored, loadErr := store.Load(t.Context(), "shared-carryover-target"); loadErr == nil || stored != nil { + t.Fatalf("rejected shared-engine session was persisted: session=%v err=%v", stored, loadErr) + } + return + } + if err != nil || created == nil { + t.Fatalf("excluded broker key blocked shared-engine create: session=%v err=%v", created, err) + } + stored, err := store.Load(t.Context(), created.ID) + if err != nil { + t.Fatal(err) + } + authority, _ := stored.BoundAuthority() + if !slices.Equal(authority.CapabilitySet.Tools, tc.authority) { + t.Fatalf("shared-engine carryover widened excluded authority: got %v, want %v", authority.CapabilitySet.Tools, tc.authority) + } + }) + } +} + +func TestLoadSessionWithMPCRetainsPriorEngineOnEnrollmentCollision(t *testing.T) { + svc, _, created := newWorkspaceEnrollmentHTTPService(t, false) + const key = "mcp__prior__tool" + replacePersistedBrokerLedger(t, svc, created, []string{key}, true) + before, err := svc.cfg.Store.Load(t.Context(), created.ID) + if err != nil { + t.Fatal(err) + } + beforeAuthority, _ := before.BoundAuthority() + beforeKeys, beforePresent := before.WorkspaceEnrollmentBrokerKeys() + svc.cfg.SessionEngine = func(context.Context, ProviderSelector, []mcp.ServerConfig, SessionProfile, string, session.PermissionMode) (SessionEngineResult, error) { + return brokerEngineResult(), nil + } + var closed bool + svc.cfg.SessionEngineWithTools = func(context.Context, ProviderSelector, []mcp.ServerConfig, SessionProfile, string, session.PermissionMode, []tool.Tool, []string) (SessionEngineResult, error) { + return SessionEngineResult{ + Engine: brokerEngineResult().Engine, + NonBrokerRegistrationKeys: []string{key}, + Close: func() error { closed = true; return nil }, + }, nil + } + svc.mu.Lock() + priorEngine := svc.sessionEngines[created.ID] + svc.mu.Unlock() + _, err = svc.LoadSessionWithMCP(t.Context(), created.ID, []mcp.ServerConfig{{Name: "client"}}) + var collision *WorkspaceEnrollmentCollisionError + if !errors.As(err, &collision) || collision.Key != key || !errors.Is(err, ErrFailedPrecondition) { + t.Fatalf("MCP load collision = %v, want typed failed-precondition %q", err, key) + } + if !closed { + t.Fatal("rejected MCP-load engine candidate was not closed") + } + svc.mu.Lock() + gotEngine := svc.sessionEngines[created.ID] + svc.mu.Unlock() + if gotEngine != priorEngine { + t.Fatal("MCP-load collision displaced the existing session engine") + } + after, err := svc.cfg.Store.Load(t.Context(), created.ID) + if err != nil { + t.Fatal(err) + } + afterAuthority, _ := after.BoundAuthority() + afterKeys, afterPresent := after.WorkspaceEnrollmentBrokerKeys() + if !reflect.DeepEqual(afterAuthority, beforeAuthority) || beforePresent != afterPresent || !slices.Equal(beforeKeys, afterKeys) { + t.Fatalf("MCP-load collision changed authority/provenance: authority=%#v keys=%q present=%t", afterAuthority, afterKeys, afterPresent) + } +} + +func TestCallSessionEngineCollisionSettlementClearsEnrollment(t *testing.T) { + svc, broker, created := newWorkspaceEnrollmentHTTPService(t, false) + const key = "mcp__prior__tool" + replacePersistedBrokerLedger(t, svc, created, []string{key}, true) + svc.cfg.SessionContextEngine = func(_ context.Context, _ session.SessionID, _ *session.Principal, _ ExecutionWorkspaceAcquirer, _ ProviderSelector, _ []mcp.ServerConfig, _ SessionProfile, _ string, _ session.PermissionMode, _ []tool.Tool, keys []string, _ session.SessionID) (SessionEngineResult, error) { + if slices.Contains(keys, "mcp__replacement__tool") { + return SessionEngineResult{}, &WorkspaceEnrollmentCollisionError{Key: key} + } + return SessionEngineResult{}, nil + } + started, err := svc.ConnectWorkspaceServices(t.Context(), created.ID) + if err != nil { + t.Fatal(err) + } + catalogue, err := brokercontract.NewWorkspaceCatalogue(started.Ref, []tool.Tool{enrollmentTool{name: "mcp__replacement__tool"}}) + if err != nil { + t.Fatal(err) + } + broker.attachment.result = brokercontract.WorkspaceEnrollmentResult{Ref: started.Ref, Status: brokercontract.WorkspaceEnrollmentConnected, Catalogue: catalogue} + _, err = svc.ConnectWorkspaceServices(t.Context(), created.ID) + var collision *WorkspaceEnrollmentCollisionError + if !errors.As(err, &collision) || collision.Key != key || !errors.Is(err, ErrFailedPrecondition) { + t.Fatalf("context factory collision = %v, want typed failed-precondition", err) + } + stored, err := svc.cfg.Store.Load(t.Context(), created.ID) + if err != nil { + t.Fatal(err) + } + if _, pending := stored.PendingWorkspaceEnrollment(); pending { + t.Fatal("context-factory collision left enrollment pending in durable session") + } +} diff --git a/internal/app/workspace_enrollment_context_test.go b/internal/app/workspace_enrollment_context_test.go new file mode 100644 index 0000000000..1c4d859d30 --- /dev/null +++ b/internal/app/workspace_enrollment_context_test.go @@ -0,0 +1,74 @@ +package app + +import ( + "context" + "errors" + "testing" + + agentfs "github.com/stacklok/mecatl/engine/adapter/agentfs" + "github.com/stacklok/mecatl/engine/adapter/memstore" + "github.com/stacklok/mecatl/engine/adapter/mockllm" + "github.com/stacklok/mecatl/engine/prompt" + "github.com/stacklok/mecatl/engine/session" + "github.com/stacklok/mecatl/engine/tool" + "github.com/stacklok/mecatl/internal/adapter/permconfig" + "github.com/stacklok/mecatl/internal/adapter/server" +) + +func TestBuildEngineContextFactoryForwardsFrozenKeysToCatalog(t *testing.T) { + kinds := harnessEmptyKinds() + kinds.Instructions = permconfig.HarnessContextKind{Sources: []string{"tenant"}, Mode: "combine"} + kinds.Commands = permconfig.HarnessContextKind{Sources: []string{"tenant"}, Mode: "combine"} + cfg := harnessPolicyConfig(t, permconfig.HarnessContextSection{EnabledSources: []string{"tenant"}, Kinds: kinds}) + cfg.HarnessInstructionSources = []HarnessSourceRegistration[prompt.InstructionAssembler]{{ + ID: "tenant", Scope: HarnessSourceScopePrincipal, Provenance: HarnessProvenancePolicy{Fixed: "driver"}, + Bind: func(context.Context, HarnessSourceScope) (prompt.InstructionAssembler, func() error, error) { + return hcAssembler("context source"), nil, nil + }, + }} + cfg.HarnessCommandSources = []HarnessSourceRegistration[server.CommandSourceBinding]{{ + ID: "tenant", Scope: HarnessSourceScopePrincipal, Provenance: HarnessProvenancePolicy{Fixed: "driver"}, + Bind: func(context.Context, HarnessSourceScope) (server.CommandSourceBinding, func() error, error) { + return &hcCommands{values: map[string]string{}}, func() error { return nil }, nil + }, + }} + cfg.permResolver = buildPermResolver(cfg) + if _, err := validateHarnessPolicy(cfg); err != nil { + t.Fatal(err) + } + prepareHarnessProcessBindings(&cfg) + if err := resolveProcessHarnessInstructions(t.Context(), &cfg); err != nil { + t.Fatal(err) + } + if err := resolveProcessHarnessSnapshots(t.Context(), &cfg); err != nil { + t.Fatal(err) + } + commandResolver, err := buildConfiguredCommandResolver(t.Context(), cfg, nil) + if err != nil { + t.Fatal(err) + } + cfg.harnessResolver = commandResolver.(*harnessCommandResolver) + if cfg.harnessContextClose != nil { + t.Cleanup(cfg.harnessContextClose) + } + + provider := mockllm.New() + reg := regForTest(provider, providerMock, "mock") + store := memstore.New() + _, _, _, _, _, _, _, assets, _, closeMCP, err := buildEngine(t.Context(), cfg, reg, provider, store, store, agentfs.NewRegistry(nil)) + if err != nil { + t.Fatal(err) + } + t.Cleanup(closeMCP) + if assets.sessionContextFactory == nil { + t.Fatal("buildEngine did not install its context-engine factory") + } + + _, err = assets.sessionContextFactory(t.Context(), "context-collision", nil, nil, server.ProviderSelector{}, nil, + server.ProfileDefault, cfg.Workspace, session.ModeDefault, + []tool.Tool{classificationNamedTool("Read")}, []string{"Read"}, "") + var collision *server.WorkspaceEnrollmentCollisionError + if !errors.As(err, &collision) || collision.Key != "Read" || !errors.Is(err, server.ErrFailedPrecondition) { + t.Fatalf("context-factory catalog collision = %v, want typed failed-precondition for Read", err) + } +}