001/*
002 * Copyright 2022-2026 Revetware LLC.
003 *
004 * Licensed under the Apache License, Version 2.0 (the "License");
005 * you may not use this file except in compliance with the License.
006 * You may obtain a copy of the License at
007 *
008 * http://www.apache.org/licenses/LICENSE-2.0
009 *
010 * Unless required by applicable law or agreed to in writing, software
011 * distributed under the License is distributed on an "AS IS" BASIS,
012 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013 * See the License for the specific language governing permissions and
014 * limitations under the License.
015 */
016
017package com.soklet;
018
019import org.jspecify.annotations.NonNull;
020
021import javax.annotation.concurrent.ThreadSafe;
022import java.util.LinkedHashMap;
023import java.util.Map;
024import java.util.Optional;
025
026import static java.util.Collections.unmodifiableMap;
027import static java.util.Objects.requireNonNull;
028
029/**
030 * Immutable key/value bag associated with an MCP session.
031 *
032 * @author <a href="https://www.revetkn.com">Mark Allen</a>
033 */
034@ThreadSafe
035public interface McpSessionContext {
036        /**
037         * Retrieves a stored session value without type conversion.
038         *
039         * @param key the session key
040         * @return the stored value, if present
041         */
042        @NonNull
043        Optional<Object> get(@NonNull String key);
044
045        /**
046         * Retrieves a stored session value with an assignability check.
047         *
048         * @param key the session key
049         * @param type the desired type
050         * @param <T> the desired type
051         * @return the stored value cast to {@code type}, if present
052         * @throws IllegalArgumentException if the stored value exists but is not assignable to {@code type}
053         */
054        @NonNull
055        <T> Optional<T> get(@NonNull String key,
056                                                                                        @NonNull Class<T> type);
057
058        /**
059         * Checks whether a session value is present for the given key.
060         *
061         * @param key the session key
062         * @return {@code true} if the key is present
063         */
064        @NonNull
065        Boolean contains(@NonNull String key);
066
067        /**
068         * Returns a new session context containing the given key/value pair.
069         *
070         * @param key the session key
071         * @param value the session value
072         * @return a new session context including the provided value
073         */
074        @NonNull
075        McpSessionContext with(@NonNull String key,
076                                                                                                 @NonNull Object value);
077
078        /**
079         * Returns a new session context without the given key.
080         *
081         * @param key the session key to remove
082         * @return a new session context without the key, or this instance if the key was absent
083         */
084        @NonNull
085        McpSessionContext without(@NonNull String key);
086
087        /**
088         * Provides an immutable snapshot of all stored session values.
089         *
090         * @return an immutable map view of the stored session values
091         */
092        @NonNull
093        Map<@NonNull String, @NonNull Object> asMap();
094
095        /**
096         * Creates an empty session context.
097         *
098         * @return a blank session context
099         */
100        @NonNull
101        static McpSessionContext fromBlankSlate() {
102                return new DefaultMcpSessionContext(Map.of());
103        }
104
105        /**
106         * Creates a session context from the provided values.
107         *
108         * @param values the values to seed into the new session context
109         * @return a new immutable session context containing {@code values}
110         */
111        @NonNull
112        static McpSessionContext fromValues(@NonNull Map<@NonNull String, @NonNull Object> values) {
113                requireNonNull(values);
114                return new DefaultMcpSessionContext(values);
115        }
116}
117
118final class DefaultMcpSessionContext implements McpSessionContext {
119        @NonNull
120        private final Map<@NonNull String, @NonNull Object> values;
121
122        DefaultMcpSessionContext(@NonNull Map<@NonNull String, @NonNull Object> values) {
123                requireNonNull(values);
124                this.values = unmodifiableMap(new LinkedHashMap<>(values));
125        }
126
127        @NonNull
128        @Override
129        public Optional<Object> get(@NonNull String key) {
130                requireNonNull(key);
131                return Optional.ofNullable(this.values.get(key));
132        }
133
134        @NonNull
135        @Override
136        public <T> Optional<T> get(@NonNull String key,
137                                                                                                                 @NonNull Class<T> type) {
138                requireNonNull(key);
139                requireNonNull(type);
140
141                Object value = this.values.get(key);
142
143                if (value == null)
144                        return Optional.empty();
145
146                if (!type.isInstance(value))
147                        throw new IllegalArgumentException("Session value for key '%s' is not assignable to %s".formatted(key, type.getName()));
148
149                return Optional.of(type.cast(value));
150        }
151
152        @NonNull
153        @Override
154        public Boolean contains(@NonNull String key) {
155                requireNonNull(key);
156                return this.values.containsKey(key);
157        }
158
159        @NonNull
160        @Override
161        public McpSessionContext with(@NonNull String key,
162                                                                                                                                @NonNull Object value) {
163                requireNonNull(key);
164                requireNonNull(value);
165
166                Map<String, Object> updatedValues = new LinkedHashMap<>(this.values);
167                updatedValues.put(key, value);
168                return new DefaultMcpSessionContext(updatedValues);
169        }
170
171        @NonNull
172        @Override
173        public McpSessionContext without(@NonNull String key) {
174                requireNonNull(key);
175
176                if (!this.values.containsKey(key))
177                        return this;
178
179                Map<String, Object> updatedValues = new LinkedHashMap<>(this.values);
180                updatedValues.remove(key);
181                return new DefaultMcpSessionContext(updatedValues);
182        }
183
184        @NonNull
185        @Override
186        public Map<@NonNull String, @NonNull Object> asMap() {
187                return this.values;
188        }
189}