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_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/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..12d0784228 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,58 @@ 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) + // 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 { + denied = true + } + } + 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) { 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/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) + } +} 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