Enable test mode on API gateway (#389)

Author: nicoloboschi

bin/langstream

@@ -37,7 +37,9 @@ if [ ! -d langstream-cli/target/cli ]; then
 fi
 popd > /dev/null
 LANGSTREAM_CLI_CONFIG=${LANGSTREAM_CLI_CONFIG:-"conf/cli.yaml"}
-echo "Using development CLI config file $(realpath $LANGSTREAM_CLI_CONFIG). To use the global config file, set LANGSTREAM_CLI_CONFIG=\$HOME/.langstream/config"
+if [ "$LANGSTREAM_CLI_CONFIG" == "conf/cli.yaml" ]; then
+  echo "Using development CLI config file $(realpath $LANGSTREAM_CLI_CONFIG). To use the global config file, set LANGSTREAM_CLI_CONFIG=\$HOME/.langstream/config"
+fi
 
 "$ROOT_DIR/langstream-cli/target/cli/bin/langstream" --conf "$LANGSTREAM_CLI_CONFIG" "$@"
 

docker/build.sh

@@ -27,11 +27,13 @@ docker_platforms() {
   fi
 }
 
+common_flags="-DskipTests -PskipPython -Dlicense.skip -Dspotless.skip"
+
 
 build_docker_image() {
   module=$1
-  ./mvnw clean install -am -DskipTests -pl $module -T 1C -PskipPython
-  ./mvnw package -DskipTests -Pdocker -pl $module -Ddocker.platforms="$(docker_platforms)" -PskipPython
+  ./mvnw install -am -pl $module -T 1C $common_flags
+  ./mvnw package -Pdocker -pl $module -Ddocker.platforms="$(docker_platforms)" $common_flags
   docker images | head -n 2
 }
 
@@ -47,9 +49,9 @@ elif [ "$only_image" == "api-gateway" ]; then
   build_docker_image langstream-api-gateway
 else
   # Build all artifacts
-  ./mvnw install -DskipTests -T 1C -Ddocker.platforms="$(docker_platforms)" -PskipPython
+  ./mvnw install -T 1C -Ddocker.platforms="$(docker_platforms)" $common_flags
   # Build docker images
-  ./mvnw package -DskipTests -Pdocker -Ddocker.platforms="$(docker_platforms)" -PskipPython
+  ./mvnw package -Pdocker -Ddocker.platforms="$(docker_platforms)" $common_flags
   docker images | head -n 6
 fi
 

examples/applications/gateway-authentication/gateways.yaml

@@ -21,20 +21,20 @@ gateways:
     topic: input-topic
     parameters:
       - sessionId
-    produceOptions:
+    produce-options:
       headers:
         - key: langstream-client-session-id
-          valueFromParameters: sessionId
+          value-from-parameters: sessionId
   - id: consume-output-no-auth
     type: consume
     topic: output-topic
     parameters:
       - sessionId
-    consumeOptions:
+    consume-options:
       filters:
         headers:
           - key: langstream-client-session-id
-            valueFromParameters: sessionId
+            value-from-parameters: sessionId
 
   - id: produce-input-auth-google
     type: produce
@@ -43,31 +43,33 @@ gateways:
       - sessionId
     authentication:
       provider: google
+      allow-test-mode: true
       configuration:
         clientId: "{{ secrets.google.client-id }}"
-    produceOptions:
+    produce-options:
       headers:
         - key: langstream-client-user-id
-          valueFromAuthentication: subject
+          value-from-authentication: subject
         - key: langstream-client-session-id
-          valueFromParameters: sessionId
+          value-from-parameters: sessionId
 
   - id: consume-output-auth-google
     type: consume
     topic: output-topic
     parameters:
       - sessionId
     authentication:
+      allow-test-mode: true
       provider: google
       configuration:
         clientId: "{{ secrets.google.client-id }}"
-    consumeOptions:
+    consume-options:
       filters:
         headers:
           - key: langstream-client-user-id
-            valueFromAuthentication: subject
+            value-from-authentication: subject
           - key: langstream-client-session-id
-            valueFromParameters: sessionId
+            value-from-parameters: sessionId
 
   - id: produce-input-auth-github
     type: produce
@@ -78,12 +80,12 @@ gateways:
       provider: github
       configuration:
         clientId: "{{ secrets.github.client-id }}"
-    produceOptions:
+    produce-options:
       headers:
         - key: langstream-client-user-id
-          valueFromAuthentication: login
+          value-from-authentication: login
         - key: langstream-client-session-id
-          valueFromParameters: sessionId
+          value-from-parameters: sessionId
 
   - id: consume-output-auth-github
     type: consume
@@ -94,10 +96,10 @@ gateways:
       provider: github
       configuration:
         clientId: "{{ secrets.github.client-id }}"
-    consumeOptions:
+    consume-options:
       filters:
         headers:
           - key: langstream-client-user-id
-            valueFromAuthentication: login
+            value-from-authentication: login
           - key: langstream-client-session-id
-            valueFromParameters: sessionId
+            value-from-parameters: sessionId

examples/applications/openai-completions/gateways.yaml

@@ -21,21 +21,21 @@ gateways:
     topic: input-topic
     parameters:
       - sessionId
-    produceOptions:
+    produce-options:
       headers:
         - key: langstream-client-session-id
-          valueFromParameters: sessionId
+          value-from-parameters: sessionId
 
   - id: consume-output
     type: consume
     topic: output-topic
     parameters:
       - sessionId
-    consumeOptions:
+    consume-options:
       filters:
         headers:
           - key: langstream-client-session-id
-            valueFromParameters: sessionId
+            value-from-parameters: sessionId
 
   - id: consume-history
     type: consume
@@ -48,10 +48,10 @@ gateways:
       provider: google
       configuration:
         clientId: "{{ secrets.google.client-id }}"
-    produceOptions:
+    produce-options:
       headers:
         - key: langstream-client-user-id
-          valueFromAuthentication: subject
+          value-from-authentication: subject
 
   - id: consume-output-auth
     type: consume
@@ -60,9 +60,9 @@ gateways:
       provider: google
       configuration:
         clientId: "{{ secrets.google.client-id }}"
-    consumeOptions:
+    consume-options:
       filters:
         headers:
           - key: langstream-client-user-id
-            valueFromAuthentication: subject
+            value-from-authentication: subject
 

langstream-api-gateway-auth/langstream-github-api-gateway-auth/src/main/java/ai/langstream/apigateway/auth/impl/github/GitHubAuthenticationProvider.java

@@ -81,7 +81,7 @@ public GatewayAuthenticationResult authenticate(GatewayRequestContext context) {
                 log.info("X-OAuth-Client-Id: {}", responseClientId);
                 log.info("Required: X-OAuth-Client-Id: {}", clientId);
 
-                Map<String, String> result = new ObjectMapper().readValue(body, Map.class);
+                Map<String, String> result = mapper.readValue(body, Map.class);
                 if (log.isDebugEnabled()) {
                     response.headers().map().forEach((k, v) -> log.debug("Header {}: {}", k, v));
                 }

langstream-api-gateway-auth/langstream-http-api-gateway-auth/pom.xml

@@ -0,0 +1,47 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!--
+
+    Copyright DataStax, Inc.
+
+    Licensed under the Apache License, Version 2.0 (the "License");
+    you may not use this file except in compliance with the License.
+    You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+    Unless required by applicable law or agreed to in writing, software
+    distributed under the License is distributed on an "AS IS" BASIS,
+    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+    See the License for the specific language governing permissions and
+    limitations under the License.
+
+-->
+<project xmlns="http://maven.apache.org/POM/4.0.0"
+         xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+  <parent>
+    <artifactId>langstream-api-gateway-auth</artifactId>
+    <groupId>ai.langstream</groupId>
+    <version>0.0.16-SNAPSHOT</version>
+  </parent>
+  <modelVersion>4.0.0</modelVersion>
+
+  <artifactId>langstream-http-api-gateway-auth</artifactId>
+
+  <dependencies>
+    <dependency>
+      <groupId>ai.langstream</groupId>
+      <artifactId>langstream-api</artifactId>
+      <version>${project.version}</version>
+    </dependency>
+    <dependency>
+      <groupId>org.slf4j</groupId>
+      <artifactId>slf4j-api</artifactId>
+      <scope>provided</scope>
+    </dependency>
+    <dependency>
+      <groupId>com.fasterxml.jackson.core</groupId>
+      <artifactId>jackson-databind</artifactId>
+    </dependency>
+  </dependencies>
+</project>

langstream-api-gateway-auth/langstream-http-api-gateway-auth/src/main/java/ai/langstream/apigateway/auth/impl/jwt/admin/HttpAuthenticationProvider.java

@@ -0,0 +1,92 @@
+/*
+ * Copyright DataStax, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package ai.langstream.apigateway.auth.impl.jwt.admin;
+
+import ai.langstream.api.gateway.GatewayAuthenticationProvider;
+import ai.langstream.api.gateway.GatewayAuthenticationResult;
+import ai.langstream.api.gateway.GatewayRequestContext;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import java.net.URI;
+import java.net.http.HttpClient;
+import java.net.http.HttpRequest;
+import java.net.http.HttpResponse;
+import java.time.Duration;
+import java.util.Map;
+import lombok.SneakyThrows;
+import lombok.extern.slf4j.Slf4j;
+
+@Slf4j
+public class HttpAuthenticationProvider implements GatewayAuthenticationProvider {
+
+    private static final ObjectMapper mapper = new ObjectMapper();
+    private HttpAuthenticationProviderConfiguration httpConfiguration;
+    private HttpClient httpClient;
+
+    @Override
+    public String type() {
+        return "http";
+    }
+
+    @Override
+    @SneakyThrows
+    public void initialize(Map<String, Object> configuration) {
+        httpConfiguration =
+                mapper.convertValue(configuration, HttpAuthenticationProviderConfiguration.class);
+        httpClient =
+                HttpClient.newBuilder()
+                        .connectTimeout(Duration.ofSeconds(30))
+                        .followRedirects(HttpClient.Redirect.ALWAYS)
+                        .build();
+    }
+
+    @Override
+    public GatewayAuthenticationResult authenticate(GatewayRequestContext context) {
+
+        final Map<String, String> placeholders = Map.of("tenant", context.tenant());
+        final String uri = resolvePlaceholders(placeholders, httpConfiguration.getPathTemplate());
+        final String url = httpConfiguration.getBaseUrl() + uri;
+
+        log.info("Authenticating admin with url: {}", url);
+
+        final HttpRequest.Builder builder = HttpRequest.newBuilder().uri(URI.create(url));
+
+        httpConfiguration.getHeaders().forEach(builder::header);
+        builder.header("Authorization", "Bearer " + context.credentials());
+        final HttpRequest request = builder.GET().build();
+
+        final HttpResponse<Void> response;
+        try {
+            response = httpClient.send(request, HttpResponse.BodyHandlers.discarding());
+        } catch (InterruptedException e) {
+            Thread.currentThread().interrupt();
+            throw new RuntimeException(e);
+        } catch (Throwable e) {
+            return GatewayAuthenticationResult.authenticationFailed(e.getMessage());
+        }
+        if (httpConfiguration.getAcceptedStatuses().contains(response.statusCode())) {
+            return GatewayAuthenticationResult.authenticationSuccessful(Map.of());
+        }
+        return GatewayAuthenticationResult.authenticationFailed(
+                "Http authentication failed: " + response.statusCode());
+    }
+
+    private static String resolvePlaceholders(Map<String, String> placeholders, String url) {
+        for (Map.Entry<String, String> entry : placeholders.entrySet()) {
+            url = url.replace("{" + entry.getKey() + "}", entry.getValue());
+        }
+        return url;
+    }
+}

langstream-api-gateway-auth/langstream-http-api-gateway-auth/src/main/java/ai/langstream/apigateway/auth/impl/jwt/admin/HttpAuthenticationProviderConfiguration.java

@@ -0,0 +1,41 @@
+/*
+ * Copyright DataStax, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package ai.langstream.apigateway.auth.impl.jwt.admin;
+
+import com.fasterxml.jackson.annotation.JsonAlias;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import lombok.AllArgsConstructor;
+import lombok.Data;
+import lombok.NoArgsConstructor;
+
+@Data
+@NoArgsConstructor
+@AllArgsConstructor
+public class HttpAuthenticationProviderConfiguration {
+
+    @JsonAlias({"base-url", "baseurl"})
+    private String baseUrl;
+
+    @JsonAlias({"path-template", "pathtemplate"})
+    private String pathTemplate;
+
+    private Map<String, String> headers = new HashMap<>();
+
+    @JsonAlias({"accepted-statuses", "acceptedstatuses"})
+    private List<Integer> acceptedStatuses = List.of(200, 201);
+}

langstream-api-gateway-auth/langstream-http-api-gateway-auth/src/main/resources/META-INF/services/ai.langstream.api.gateway.GatewayAuthenticationProvider

@@ -0,0 +1 @@
+ai.langstream.apigateway.auth.impl.jwt.admin.HttpAuthenticationProvider