feat: organization session policy SDK methods (#152)

* feat: add organization session policy SDK methods and generate local proto target

Adds get_organization_session_policy and update_organization_session_policy
wrapper methods, regenerates gRPC stubs from local protos, and adds
generate-local Makefile target.

* chore: bump proto ref to v0.1.120.2

* feat: regenerate protos for v0.1.120.2 and add session policy tests

Updates generated stubs to match new flat UpdateOrganizationSessionPolicyRequest
(SessionPolicyType enum rename, flat fields). Adds agentkit_logs generated module
and removes stale buf/validate generated files. Adds integration test coverage
for get/update/revert organization session policy.

* fix: add --include-imports to generate-local and restore buf/validate stubs

The generate-local target was missing --include-imports, causing buf/validate
generated files to be absent while generated stubs still imported them.

* chore: revert unrelated proto changes, remove generate-local, bump PROTO_REF

* chore: ignore proto/ directory (local generate-local output)

* chore: regenerate all proto stubs for v0.1.120.2

* chore: regenerate proto stubs from local scalekit proto and add generate-local target

Resets pb2 files to main baseline then regenerates using local ../scalekit proto
via `buf generate ../scalekit --include-imports` to avoid GitHub clone and resolve
buf/validate deps correctly. Adds generate-local Makefile target for future use.

* build

* fix: remove blocking input() call from test_actions setUp

* Revert "fix: remove blocking input() call from test_actions setUp"

This reverts commit 8c84f5b.

* build

* chore: regenerate proto stubs from local scalekit proto

* chore: regenerate proto stubs from published v0.1.121.2

* fix: address coderabbitai review comments

- generate-local: rename trap handler to rollback_and_cleanup, add cleanup
  of TEMP_DIR and intermediate dirs on failure, and clear trap on success
- test_revert_to_application_policy: add follow-up get to verify persisted state

* build

* fix: return raw grpc response from session policy methods and bump version to 2.10.0

- Remove response[0].policy extraction from get/update_organization_session_policy to match existing SDK patterns
- Bump SDK version 2.9.0 → 2.10.0
- Update api_version date to 20260513

* fix: update session policy tests to unpack raw grpc response tuple

* fix: correct return types for session policy methods to use response types

Author: srinivaskarre-sk

.gitignore

@@ -170,7 +170,8 @@ proto
 
 # Claude Code documentation
 CLAUDE.md
+proto/
 .claude/*
 
 # macOS
-.DS_Store
+.DS_Store

Makefile

@@ -14,6 +14,7 @@ VENV_PIP := $(VENV_PYTHON) -m pip
 PROTO_REPO_URL := https://github.com/scalekit-inc/scalekit.git
 PROTO_REF ?= v0.1.121.2
 PROTO_SUBDIR := proto
+LOCAL_PROTO_REPO ?= ../scalekit
 
 TEMP_DIR := temp_scalekit
 SCALEKIT_DIR := scalekit
@@ -22,7 +23,7 @@ GOOGLE_DIR := google
 PROTO_DIR := proto
 PROTOC_DIR := protoc_gen_openapiv2
 
-.PHONY: setup generate lint test tools-check create-venv prepare buf_generate restore generate_init_files cleanup copy_proto_dir
+.PHONY: setup generate generate-local lint test tools-check create-venv prepare buf_generate restore generate_init_files cleanup copy_proto_dir
 
 setup: create-venv
 	@echo "Installing SDK dependencies in $(VENV_DIR)..."
@@ -109,6 +110,28 @@ cleanup:
 	rm -f .dirpath buf.yaml buf.lock
 	@if [ -f buf.work.yaml.bak ]; then mv buf.work.yaml.bak buf.work.yaml; fi
 
+generate-local: tools-check
+	@echo "Using local proto sources from $(LOCAL_PROTO_REPO)..."
+	@set -euo pipefail; \
+	prepared=0; \
+	rollback_and_cleanup() { \
+		if [ "$$prepared" -eq 1 ] && [ -d "$(TEMP_DIR)" ]; then \
+			echo "Generation failed; restoring $(SCALEKIT_DIR) from $(TEMP_DIR)..."; \
+			rsync -a "$(TEMP_DIR)/" "$(SCALEKIT_DIR)/"; \
+		fi; \
+		rm -rf "$(TEMP_DIR)" "$(GOOGLE_DIR)" "$(PROTO_DIR)" "$(PROTOC_DIR)"; \
+		rm -f .dirpath buf.yaml buf.lock; \
+		if [ -f buf.work.yaml.bak ]; then mv buf.work.yaml.bak buf.work.yaml; fi; \
+	}; \
+	trap 'rollback_and_cleanup' EXIT; \
+	$(MAKE) prepare; prepared=1; \
+	buf generate $(LOCAL_PROTO_REPO) --include-imports; \
+	$(MAKE) restore; prepared=0; \
+	$(MAKE) generate_init_files; \
+	$(MAKE) cleanup; \
+	trap - EXIT
+	@echo "Code generation complete."
+
 lint: create-venv
 	@echo "Running static checks..."
 	$(VENV_PYTHON) -m compileall -q scalekit tests

scalekit/_version.py

@@ -1,3 +1,3 @@
 # Single source of truth for the SDK version.
 # Import this in setup.py and scalekit/core.py — never hardcode the version elsewhere.
-__version__ = "2.9.0"
+__version__ = "2.10.0"

scalekit/core.py

@@ -29,7 +29,7 @@ class CoreClient:
 
     sdk_version = f"Scalekit-Python/{_sdk_version}"
     # YYYYMMDD
-    api_version = "20260428"
+    api_version = "20260513"
     user_agent = f"{sdk_version} Python/{platform.python_version()} ({platform.system()}; {platform.architecture()}"
 
     def __init__(self, env_url, client_id, client_secret):

scalekit/organization.py

@@ -20,7 +20,13 @@
     UpdateOrganizationSettingsRequest,
     OrganizationUserManagementSettings,
     UpsertUserManagementSettingsRequest,
+    GetOrganizationSessionPolicyRequest,
+    GetOrganizationSessionPolicyResponse,
+    UpdateOrganizationSessionPolicyRequest,
+    UpdateOrganizationSessionPolicyResponse,
+    SessionPolicyType,
 )
+from scalekit.v1.commons.commons_pb2 import TimeUnit
 from scalekit.v1.organizations.organizations_pb2_grpc import OrganizationServiceStub
 
 
@@ -219,3 +225,72 @@ def upsert_user_management_settings(self, organization_id: str, max_allowed_user
             )
         )
         return response[0].settings
+
+    def get_organization_session_policy(self, organization_id: str) -> GetOrganizationSessionPolicyResponse:
+        """
+        Get the session policy for an organization.
+
+        :param organization_id: Organization id
+        :type organization_id : ``` str ```
+        :returns:
+            GetOrganizationSessionPolicyResponse
+        """
+        return self.core_client.grpc_exec(
+            self.organization_service.GetOrganizationSessionPolicy.with_call,
+            GetOrganizationSessionPolicyRequest(organization_id=organization_id),
+        )
+
+    def update_organization_session_policy(
+        self,
+        organization_id: str,
+        policy_source: SessionPolicyType,
+        absolute_session_timeout: Optional[int] = None,
+        absolute_session_timeout_unit: Optional[TimeUnit] = None,
+        idle_session_timeout_enabled: Optional[bool] = None,
+        idle_session_timeout: Optional[int] = None,
+        idle_session_timeout_unit: Optional[TimeUnit] = None,
+    ) -> UpdateOrganizationSessionPolicyResponse:
+        """
+        Set a custom session policy for an organization or revert to application defaults.
+
+        :param organization_id: Organization id
+        :type organization_id: ``` str ```
+        :param policy_source: SessionPolicyType.APPLICATION or SessionPolicyType.CUSTOM
+        :type policy_source: ``` SessionPolicyType ```
+        :param absolute_session_timeout: Absolute session timeout value (optional)
+        :type absolute_session_timeout: ``` int | None ```
+        :param absolute_session_timeout_unit: Unit for absolute timeout (optional)
+        :type absolute_session_timeout_unit: ``` TimeUnit | None ```
+        :param idle_session_timeout_enabled: Whether idle session timeout is enabled (optional)
+        :type idle_session_timeout_enabled: ``` bool | None ```
+        :param idle_session_timeout: Idle session timeout value (optional)
+        :type idle_session_timeout: ``` int | None ```
+        :param idle_session_timeout_unit: Unit for idle timeout (optional)
+        :type idle_session_timeout_unit: ``` TimeUnit | None ```
+        :returns:
+            UpdateOrganizationSessionPolicyResponse
+        """
+        req = UpdateOrganizationSessionPolicyRequest(
+            organization_id=organization_id,
+            policy_source=policy_source,
+        )
+        if absolute_session_timeout is not None:
+            req.absolute_session_timeout.CopyFrom(
+                wrappers_pb2.Int32Value(value=absolute_session_timeout)
+            )
+        if absolute_session_timeout_unit is not None:
+            req.absolute_session_timeout_unit = absolute_session_timeout_unit
+        if idle_session_timeout_enabled is not None:
+            req.idle_session_timeout_enabled.CopyFrom(
+                wrappers_pb2.BoolValue(value=idle_session_timeout_enabled)
+            )
+        if idle_session_timeout is not None:
+            req.idle_session_timeout.CopyFrom(
+                wrappers_pb2.Int32Value(value=idle_session_timeout)
+            )
+        if idle_session_timeout_unit is not None:
+            req.idle_session_timeout_unit = idle_session_timeout_unit
+        return self.core_client.grpc_exec(
+            self.organization_service.UpdateOrganizationSessionPolicy.with_call,
+            req,
+        )

tests/test_organization_session_policy.py

@@ -0,0 +1,107 @@
+from faker import Faker
+
+from basetest import BaseTest
+from scalekit.v1.organizations.organizations_pb2 import (
+    CreateOrganization,
+    SessionPolicyType,
+)
+from scalekit.v1.commons.commons_pb2 import TimeUnit
+
+
+class TestOrganizationSessionPolicy(BaseTest):
+    """Test cases for organization session policy management."""
+
+    def setUp(self):
+        self.org_id = None
+
+    def _create_org(self):
+        organization = CreateOrganization(display_name=Faker().company(), external_id=Faker().uuid4())
+        response = self.scalekit_client.organization.create_organization(organization=organization)
+        self.org_id = response[0].organization.id
+        return self.org_id
+
+    def test_get_default_policy(self):
+        """New org should inherit APPLICATION policy by default."""
+        org_id = self._create_org()
+
+        policy = self.scalekit_client.organization.get_organization_session_policy(
+            organization_id=org_id
+        )[0].policy
+
+        self.assertIsNotNone(policy)
+        self.assertEqual(policy.policy_source, SessionPolicyType.APPLICATION)
+
+    def test_set_custom_policy(self):
+        """Setting a custom policy should persist and be retrievable."""
+        org_id = self._create_org()
+
+        policy = self.scalekit_client.organization.update_organization_session_policy(
+            organization_id=org_id,
+            policy_source=SessionPolicyType.CUSTOM,
+            absolute_session_timeout=360,
+            absolute_session_timeout_unit=TimeUnit.MINUTES,
+            idle_session_timeout_enabled=True,
+            idle_session_timeout=60,
+            idle_session_timeout_unit=TimeUnit.MINUTES,
+        )[0].policy
+
+        self.assertIsNotNone(policy)
+        self.assertEqual(policy.policy_source, SessionPolicyType.CUSTOM)
+
+        fetched = self.scalekit_client.organization.get_organization_session_policy(
+            organization_id=org_id
+        )[0].policy
+        self.assertEqual(fetched.policy_source, SessionPolicyType.CUSTOM)
+        self.assertTrue(fetched.HasField("absolute_session_timeout"))
+        self.assertEqual(fetched.absolute_session_timeout.value, 360)
+        self.assertTrue(fetched.HasField("idle_session_timeout_enabled"))
+        self.assertTrue(fetched.idle_session_timeout_enabled.value)
+
+    def test_revert_to_application_policy(self):
+        """Setting policy source to APPLICATION should revert to application defaults."""
+        org_id = self._create_org()
+
+        self.scalekit_client.organization.update_organization_session_policy(
+            organization_id=org_id,
+            policy_source=SessionPolicyType.CUSTOM,
+            absolute_session_timeout=120,
+            absolute_session_timeout_unit=TimeUnit.MINUTES,
+        )
+
+        reverted = self.scalekit_client.organization.update_organization_session_policy(
+            organization_id=org_id,
+            policy_source=SessionPolicyType.APPLICATION,
+        )[0].policy
+
+        self.assertIsNotNone(reverted)
+        self.assertEqual(reverted.policy_source, SessionPolicyType.APPLICATION)
+
+        fetched = self.scalekit_client.organization.get_organization_session_policy(
+            organization_id=org_id
+        )[0].policy
+        self.assertEqual(fetched.policy_source, SessionPolicyType.APPLICATION)
+
+    def test_set_idle_timeout_disabled(self):
+        """Setting idle_session_timeout_enabled=False should persist as false."""
+        org_id = self._create_org()
+
+        policy = self.scalekit_client.organization.update_organization_session_policy(
+            organization_id=org_id,
+            policy_source=SessionPolicyType.CUSTOM,
+            absolute_session_timeout=480,
+            absolute_session_timeout_unit=TimeUnit.MINUTES,
+            idle_session_timeout_enabled=False,
+        )[0].policy
+
+        self.assertIsNotNone(policy)
+        self.assertEqual(policy.policy_source, SessionPolicyType.CUSTOM)
+
+        fetched = self.scalekit_client.organization.get_organization_session_policy(
+            organization_id=org_id
+        )[0].policy
+        self.assertTrue(fetched.HasField("idle_session_timeout_enabled"))
+        self.assertFalse(fetched.idle_session_timeout_enabled.value)
+
+    def tearDown(self):
+        if self.org_id:
+            self.scalekit_client.organization.delete_organization(organization_id=self.org_id)