fix: resolve stale connection resets and unretried DEADLINE_EXCEEDED errors

Makefile

@@ -14,7 +14,7 @@ MVN_CMD := $(MVN) $(MVN_FLAGS)
 BUF := PATH="$(TOOLS_BIN):$$PATH" buf
 
 PROTO_REPO_URL := https://github.com/scalekit-inc/scalekit.git
-PROTO_REF := v0.1.103
+PROTO_REF := v0.1.121.2
 PROTO_SUBDIR := proto
 PROTO_REMOTE_INPUT := $(PROTO_REPO_URL)\#ref=$(PROTO_REF),subdir=$(PROTO_SUBDIR)
 BUF_GENERATE_FLAGS := --include-imports

src/main/java/com/scalekit/api/OrganizationClient.java

@@ -30,4 +30,14 @@ public interface OrganizationClient {
 
   OrganizationUserManagementSettings upsertUserManagementSettings(String organizationId, OrganizationUserManagementSettings settings);
 
+  GetOrganizationUserManagementSettingsResponse getOrganizationUserManagementSetting(String organizationId);
+
+  SearchOrganizationsResponse searchOrganizations(SearchOrganizationsRequest request);
+
+  UpdateOrganizationSessionPolicyResponse updateOrganizationSessionPolicy(String organizationId, UpdateOrganizationSessionPolicyRequest request);
+
+  GetOrganizationSessionPolicyResponse getOrganizationSessionPolicy(String organizationId);
+
+  GetApplicationSessionPolicyResponse getApplicationSessionPolicy(String organizationId);
+
 }

src/main/java/com/scalekit/api/UserClient.java

@@ -30,4 +30,8 @@ public interface UserClient {
     ListUserRolesResponse listUserRoles(String organizationId, String userId);
 
     ListUserPermissionsResponse listUserPermissions(String organizationId, String userId);
-} 
+
+    AssignUserRolesResponse assignUserRoles(String organizationId, String userId, AssignUserRolesRequest request);
+
+    void removeUserRole(String organizationId, String userId, RemoveUserRoleRequest request);
+}

src/main/java/com/scalekit/api/impl/ScalekitAuthClient.java

@@ -27,6 +27,7 @@
 import java.io.IOException;
 import java.io.UnsupportedEncodingException;
 import java.net.*;
+import java.util.concurrent.TimeUnit;
 
 import java.nio.charset.StandardCharsets;
 import java.util.*;
@@ -48,9 +49,14 @@ public ScalekitAuthClient() {
                 .setConnectionRequestTimeout(5000)
                 .build();
 
-        // Create an HttpClient with the custom request configuration
-         this.httpClient = HttpClients.custom()
+        // Create an HttpClient with the custom request configuration.
+        // evictExpiredConnections + evictIdleConnections prevent stale pooled connections
+        // from being reused after the server has closed them (manifests as Connection reset
+        // on token refresh when the access token has been valid for a long time).
+        this.httpClient = HttpClients.custom()
                 .setDefaultRequestConfig(requestConfig)
+                .evictExpiredConnections()
+                .evictIdleConnections(30, TimeUnit.SECONDS)
                 .build();
 
         this.objectMapper = new ObjectMapper();

src/main/java/com/scalekit/api/impl/ScalekitOrganizationClient.java

@@ -282,5 +282,92 @@ public OrganizationUserManagementSettings upsertUserManagementSettings(String or
         }, this.credentials);
     }
 
+    /**
+     * getOrganizationUserManagementSetting retrieves user management settings for an organization.
+     * @param organizationId Organization identifier
+     * @return GetOrganizationUserManagementSettingsResponse containing the settings
+     */
+    @Override
+    public GetOrganizationUserManagementSettingsResponse getOrganizationUserManagementSetting(String organizationId) {
+        return RetryExecuter.executeWithRetry(() -> {
+            return this.organizationStub
+                    .withDeadlineAfter(Environment.defaultConfig().timeout, TimeUnit.MILLISECONDS)
+                    .getOrganizationUserManagementSetting(
+                            GetOrganizationUserManagementSettingsRequest.newBuilder()
+                                    .setOrganizationId(organizationId)
+                                    .build()
+                    );
+        }, this.credentials);
+    }
+
+    /**
+     * searchOrganizations searches organizations using the provided request criteria.
+     * @param request SearchOrganizationsRequest containing query and pagination options
+     * @return SearchOrganizationsResponse containing matching organizations
+     */
+    @Override
+    public SearchOrganizationsResponse searchOrganizations(SearchOrganizationsRequest request) {
+        return RetryExecuter.executeWithRetry(() -> {
+            return this.organizationStub
+                    .withDeadlineAfter(Environment.defaultConfig().timeout, TimeUnit.MILLISECONDS)
+                    .searchOrganization(request);
+        }, this.credentials);
+    }
+
+    /**
+     * updateOrganizationSessionPolicy updates the session policy for an organization.
+     * @param organizationId Organization identifier
+     * @param request UpdateOrganizationSessionPolicyRequest containing the policy details
+     * @return UpdateOrganizationSessionPolicyResponse with the updated policy
+     */
+    @Override
+    public UpdateOrganizationSessionPolicyResponse updateOrganizationSessionPolicy(String organizationId, UpdateOrganizationSessionPolicyRequest request) {
+        return RetryExecuter.executeWithRetry(() -> {
+            return this.organizationStub
+                    .withDeadlineAfter(Environment.defaultConfig().timeout, TimeUnit.MILLISECONDS)
+                    .updateOrganizationSessionPolicy(
+                            request.toBuilder()
+                                    .setOrganizationId(organizationId)
+                                    .build()
+                    );
+        }, this.credentials);
+    }
+
+    /**
+     * getOrganizationSessionPolicy retrieves the session policy for an organization.
+     * @param organizationId Organization identifier
+     * @return GetOrganizationSessionPolicyResponse with the current policy
+     */
+    @Override
+    public GetOrganizationSessionPolicyResponse getOrganizationSessionPolicy(String organizationId) {
+        return RetryExecuter.executeWithRetry(() -> {
+            return this.organizationStub
+                    .withDeadlineAfter(Environment.defaultConfig().timeout, TimeUnit.MILLISECONDS)
+                    .getOrganizationSessionPolicy(
+                            GetOrganizationSessionPolicyRequest.newBuilder()
+                                    .setOrganizationId(organizationId)
+                                    .build()
+                    );
+        }, this.credentials);
+    }
+
+    /**
+     * getApplicationSessionPolicy retrieves the application-level session policy for an organization.
+     * @param organizationId Organization identifier
+     * @return GetApplicationSessionPolicyResponse with the application session policy
+     */
+    @Override
+    public GetApplicationSessionPolicyResponse getApplicationSessionPolicy(String organizationId) {
+        return RetryExecuter.executeWithRetry(() -> {
+            return this.organizationStub
+                    .withDeadlineAfter(Environment.defaultConfig().timeout, TimeUnit.MILLISECONDS)
+                    .getApplicationSessionPolicy(
+                            GetApplicationSessionPolicyRequest.newBuilder()
+                                    .setOrganizationId(organizationId)
+                                    .build()
+                    );
+        }, this.credentials);
+    }
+
 
 }

src/main/java/com/scalekit/api/impl/ScalekitUserClient.java

@@ -279,4 +279,42 @@ public ListUserPermissionsResponse listUserPermissions(String organizationId, St
                     .listUserPermissions(request);
         }, this.credentials);
     }
-} 
+
+    /**
+     * Assigns roles to a user in the specified organization
+     * @param organizationId: The organization ID
+     * @param userId: The ID of the user to assign roles to
+     * @param request: The assign user roles request containing the roles to assign
+     * @return AssignUserRolesResponse: The response containing the assigned roles
+     */
+    @Override
+    public AssignUserRolesResponse assignUserRoles(String organizationId, String userId, AssignUserRolesRequest request) {
+        return RetryExecuter.executeWithRetry(() -> {
+            return userService
+                    .withDeadlineAfter(Environment.defaultConfig().timeout, TimeUnit.MILLISECONDS)
+                    .assignUserRoles(request.toBuilder()
+                            .setOrganizationId(organizationId)
+                            .setUserId(userId)
+                            .build());
+        }, this.credentials);
+    }
+
+    /**
+     * Removes a role from a user in the specified organization
+     * @param organizationId: The organization ID
+     * @param userId: The ID of the user to remove the role from
+     * @param request: The remove user role request containing the role to remove
+     */
+    @Override
+    public void removeUserRole(String organizationId, String userId, RemoveUserRoleRequest request) {
+        RetryExecuter.executeWithRetry(() -> {
+            userService
+                    .withDeadlineAfter(Environment.defaultConfig().timeout, TimeUnit.MILLISECONDS)
+                    .removeUserRole(request.toBuilder()
+                            .setOrganizationId(organizationId)
+                            .setUserId(userId)
+                            .build());
+            return null;
+        }, this.credentials);
+    }
+}