Merge remote-tracking branch origin/master into deploy/prod-telemetry

.github/workflows/deploy-prod.yml

@@ -15,7 +15,7 @@ concurrency:
 env:
   CORE_IMAGE: ${{ secrets.DOCKER_USERNAME }}/langbot
   CLOUD_IMAGE: ${{ secrets.DOCKER_USERNAME }}/langbot-cloud-core
-  SPACE_REF: 178b1634c104e635c706e1fb4ca37cb3de073cf9
+  SPACE_REF: 58253c53933f95d81b035fbe2efedb55b6c1a82b
 
 jobs:
   build-and-deploy:

src/langbot/pkg/api/http/controller/groups/user.py

@@ -1,6 +1,7 @@
 import quart
 import argon2
 import asyncio
+import datetime
 import uuid
 from urllib.parse import parse_qs, urlsplit
 
@@ -218,7 +219,22 @@ class UserRouterGroup(group.RouterGroup):
             try:
                 consumed_state = await self.ap.user_service.consume_space_oauth_state_details(state, 'login')
                 # Exchange code for tokens
-                token_data = await self.ap.space_service.exchange_oauth_code(code)
+                launch_workspace_uuid = consumed_state.launch_workspace_uuid
+                workspace_uuids = [launch_workspace_uuid] if launch_workspace_uuid else []
+                workspace_created_ats: dict[str, int] = {}
+                if not workspace_uuids and getattr(getattr(self.ap, 'deployment', None), 'mode', 'oss') != 'cloud':
+                    binding = await self.ap.workspace_service.get_execution_binding()
+                    workspace_uuids = [binding.workspace_uuid]
+                    workspace_created_at = binding.workspace_created_at
+                    if workspace_created_at is not None:
+                        if workspace_created_at.tzinfo is None:
+                            workspace_created_at = workspace_created_at.replace(tzinfo=datetime.UTC)
+                        workspace_created_ats[binding.workspace_uuid] = int(workspace_created_at.timestamp())
+                token_data = await self.ap.space_service.exchange_oauth_code(
+                    code,
+                    workspace_uuids,
+                    workspace_created_ats,
+                )
                 access_token = token_data.get('access_token')
                 refresh_token = token_data.get('refresh_token')
                 expires_in = token_data.get('expires_in', 0)
@@ -231,7 +247,6 @@ class UserRouterGroup(group.RouterGroup):
                     access_token, refresh_token, expires_in
                 )
 
-                launch_workspace_uuid = consumed_state.launch_workspace_uuid
                 if launch_workspace_uuid:
                     try:
                         access = await self.ap.workspace_collaboration_service.resolve_account_workspace(

src/langbot/pkg/api/http/service/space.py

@@ -59,6 +59,10 @@ class SpaceService:
         result_list = result.all()
         return result_list[0] if result_list else None
 
+    async def get_valid_access_token(self, user_email: str) -> str | None:
+        """Return a current Space bearer, refreshing and persisting it when needed."""
+        return await self._ensure_valid_token(user_email)
+
     async def _ensure_valid_token(self, user_email: str) -> str | None:
         """Ensure access token is valid, refresh if expired. Returns valid access_token or None."""
         user_obj = await self._get_user_by_email(user_email)
@@ -117,7 +121,12 @@ class SpaceService:
             params['state'] = state
         return f'{authorize_url}?{urlencode(params)}'
 
-    async def exchange_oauth_code(self, code: str) -> typing.Dict:
+    async def exchange_oauth_code(
+        self,
+        code: str,
+        workspace_uuids: list[str] | None = None,
+        workspace_created_ats: dict[str, int] | None = None,
+    ) -> typing.Dict:
         """Exchange OAuth authorization code for tokens"""
         from langbot.pkg.utils import constants
 
@@ -127,7 +136,14 @@ class SpaceService:
         session = httpclient.get_session()
         async with session.post(
             f'{space_url}/api/v1/accounts/oauth/token',
-            json={'code': code, 'instance_id': constants.instance_id},
+            json={
+                'code': code,
+                'instance_id': constants.instance_id,
+                # Sending an explicit empty list tells new Space servers not to
+                # synthesize a legacy instance-derived Workspace binding.
+                'workspace_uuids': workspace_uuids if workspace_uuids is not None else [],
+                'workspace_created_ats': workspace_created_ats or {},
+            },
         ) as response:
             if response.status != 200:
                 error = await httpclient.read_text_limited(response)

src/langbot/pkg/api/http/service/user.py

@@ -779,8 +779,27 @@ class UserService:
         local_account = await self.get_user_by_email(user_email)
         if local_account is None:
             raise ValueError('User not found')
-        # Exchange code for tokens
-        token_data = await self.ap.space_service.exchange_oauth_code(code)
+        # Exchange code for tokens and bind both installation and the active
+        # OSS Workspace as independent identities.
+        workspace_service = getattr(self.ap, 'workspace_service', None)
+        if workspace_service is not None:
+            binding = await workspace_service.get_execution_binding()
+            created_at = binding.workspace_created_at
+            created_ts = (
+                int(created_at.replace(tzinfo=datetime.timezone.utc).timestamp())
+                if created_at.tzinfo is None
+                else int(created_at.timestamp())
+            )
+            token_data = await self.ap.space_service.exchange_oauth_code(
+                code,
+                [binding.workspace_uuid],
+                {binding.workspace_uuid: created_ts},
+            )
+        else:
+            # Compatibility for early/bootstrap call sites that have not wired
+            # WorkspaceService yet; old Space servers still derive the legacy
+            # Workspace identity from instance_id when the field is omitted.
+            token_data = await self.ap.space_service.exchange_oauth_code(code)
         access_token = token_data.get('access_token')
         refresh_token = token_data.get('refresh_token')
         expires_in = token_data.get('expires_in', 0)

src/langbot/pkg/telemetry/heartbeat.py

@@ -120,6 +120,7 @@ async def build_heartbeat_payload(
     ap: core_app.Application,
     *,
     workspace_uuid: str,
+    workspace_create_ts: int = 0,
     workspace_resource: WorkspaceResourceSnapshot | None = None,
 ) -> dict:
     """Collect one anonymous Workspace profile snapshot."""
@@ -212,7 +213,9 @@ async def build_heartbeat_payload(
         'event_type': 'instance_heartbeat',
         'query_id': '',
         'version': constants.semantic_version,
+        'instance_id': constants.instance_id,
         'workspace_uuid': workspace_uuid,
+        'workspace_create_ts': workspace_create_ts,
         'instance_create_ts': constants.instance_create_ts,
         'edition': constants.edition,
         'features': features,
@@ -220,10 +223,24 @@ async def build_heartbeat_payload(
     }
 
 
+def _workspace_created_timestamp(created_at: datetime | None) -> int:
+    if created_at is None:
+        return 0
+    if created_at.tzinfo is None:
+        # SQLAlchemy may return persisted UTC values without tzinfo. Never
+        # reinterpret them in the host's local timezone.
+        created_at = created_at.replace(tzinfo=timezone.utc)
+    return int(created_at.timestamp())
+
+
 async def build_heartbeat_payloads(ap: core_app.Application) -> list[dict]:
     """Build one heartbeat per active Workspace."""
     bindings = await ap.workspace_service.list_active_execution_bindings()
     workspace_uuids = sorted({binding.workspace_uuid for binding in bindings})
+    workspace_create_ts = {
+        binding.workspace_uuid: _workspace_created_timestamp(getattr(binding, 'workspace_created_at', None))
+        for binding in bindings
+    }
     resources = {
         resource['workspace_uuid']: resource for resource in await _cloud_workspace_resource_counts(ap, bindings)
     }
@@ -231,6 +248,7 @@ async def build_heartbeat_payloads(ap: core_app.Application) -> list[dict]:
         await build_heartbeat_payload(
             ap,
             workspace_uuid=workspace_uuid,
+            workspace_create_ts=workspace_create_ts.get(workspace_uuid, 0),
             workspace_resource=resources.get(workspace_uuid),
         )
         for workspace_uuid in workspace_uuids

src/langbot/pkg/telemetry/identity.py

@@ -4,13 +4,19 @@ import typing
 
 
 class WorkspaceExecutionContext(typing.Protocol):
+    @property
+    def instance_uuid(self) -> str: ...
+
     @property
     def workspace_uuid(self) -> str: ...
 
 
 def workspace_identity(execution_context: WorkspaceExecutionContext) -> dict[str, str]:
-    """Build the canonical telemetry identity for one Workspace execution."""
+    """Build both first-class telemetry identities for one execution."""
+    instance_id = execution_context.instance_uuid.strip()
     workspace_uuid = execution_context.workspace_uuid.strip()
+    if not instance_id:
+        raise ValueError('Telemetry execution instance ID is empty')
     if not workspace_uuid:
         raise ValueError('Telemetry execution Workspace UUID is empty')
-    return {'workspace_uuid': workspace_uuid}
+    return {'instance_id': instance_id, 'workspace_uuid': workspace_uuid}

src/langbot/pkg/telemetry/telemetry.py

@@ -136,12 +136,31 @@ class TelemetryManager:
                     try:
                         # Use asyncio.wait_for to ensure we always bound the total time
                         telemetry_token = os.getenv('LANGBOT_TELEMETRY_INGEST_TOKEN', '').strip()
+                        headers: dict[str, str] = {}
                         if telemetry_token:
-                            request = client.post(
-                                url,
-                                json=sanitized,
-                                headers={'X-LangBot-Telemetry-Token': telemetry_token},
-                            )
+                            headers['X-LangBot-Telemetry-Token'] = telemetry_token
+                        else:
+                            workspace_uuid = str(sanitized.get('workspace_uuid', '')).strip()
+                            user_service = getattr(self.ap, 'user_service', None)
+                            if workspace_uuid and user_service is not None:
+                                try:
+                                    owner = await user_service.get_workspace_owner(workspace_uuid)
+                                    owner_email = str(getattr(owner, 'user', '') or '').strip()
+                                    space_service = getattr(self.ap, 'space_service', None)
+                                    access_token = (
+                                        await space_service.get_valid_access_token(owner_email)
+                                        if owner_email and space_service is not None
+                                        else None
+                                    )
+                                    access_token = str(access_token or '').strip()
+                                    if access_token:
+                                        headers['Authorization'] = f'Bearer {access_token}'
+                                except Exception:
+                                    self.ap.logger.debug(
+                                        'Could not resolve authenticated telemetry reporter', exc_info=True
+                                    )
+                        if headers:
+                            request = client.post(url, json=sanitized, headers=headers)
                         else:
                             request = client.post(url, json=sanitized)
                         resp = await asyncio.wait_for(request, timeout=10 + 1)

src/langbot/pkg/workspace/entities.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import datetime
 from dataclasses import dataclass
 
 
@@ -12,3 +13,4 @@ class WorkspaceExecutionBinding:
     placement_generation: int
     write_fenced: bool
     state: str
+    workspace_created_at: datetime.datetime | None = None

src/langbot/pkg/workspace/service.py

@@ -283,6 +283,7 @@ class WorkspaceService:
                 placement_generation=execution_state.active_generation,
                 write_fenced=execution_state.write_fenced,
                 state=execution_state.state,
+                workspace_created_at=workspace.created_at,
             )
 
         binding = await self._run(operation, session=session)

tests/integration/api/test_user_space_oauth.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import datetime
 from types import SimpleNamespace
 from unittest.mock import AsyncMock, Mock
 from urllib.parse import parse_qs, urlsplit
@@ -14,6 +15,7 @@ from langbot.pkg.api.http.controller.groups.user import UserRouterGroup
 
 pytestmark = pytest.mark.integration
 WORKSPACE_UUID = '11111111-1111-4111-8111-111111111111'
+WORKSPACE_CREATED_AT = datetime.datetime(2026, 1, 2, 3, 4, 5, tzinfo=datetime.UTC)
 
 
 @pytest.fixture
@@ -58,6 +60,12 @@ async def space_oauth_api():
         return_value={'account_uuid': 'account-a', 'workspace_uuid': WORKSPACE_UUID}
     )
     application.workspace_collaboration_service.resolve_account_workspace = AsyncMock(return_value=access)
+    application.workspace_service.get_execution_binding = AsyncMock(
+        return_value=SimpleNamespace(
+            workspace_uuid=WORKSPACE_UUID,
+            workspace_created_at=WORKSPACE_CREATED_AT,
+        )
+    )
     application.space_service.get_oauth_authorize_url = Mock(
         side_effect=lambda redirect_uri, state: f'https://space.example/authorize?state={state}'
     )
@@ -234,7 +242,11 @@ async def test_login_callback_requires_and_consumes_server_state(space_oauth_api
     assert response.status_code == 200
     assert (await response.get_json())['data']['token'] == 'space-login-token'
     application.user_service.consume_space_oauth_state_details.assert_awaited_once_with('opaque-login-state', 'login')
-    application.space_service.exchange_oauth_code.assert_awaited_once_with('oauth-code')
+    application.space_service.exchange_oauth_code.assert_awaited_once_with(
+        'oauth-code',
+        [WORKSPACE_UUID],
+        {WORKSPACE_UUID: int(WORKSPACE_CREATED_AT.timestamp())},
+    )
 
 
 @pytest.mark.asyncio

tests/unit_tests/api/service/test_space_service.py

@@ -25,6 +25,7 @@ import time
 
 from langbot.pkg.api.http.service.space import SpaceService
 from langbot.pkg.entity.persistence.user import User
+from langbot.pkg.utils import constants
 
 
 pytestmark = pytest.mark.asyncio
@@ -573,10 +574,20 @@ class TestSpaceServiceExchangeOAuthCode:
             mock_session_obj.post.return_value.__aexit__ = AsyncMock(return_value=None)
 
             # Execute
-            result = await service.exchange_oauth_code('auth_code')
+            result = await service.exchange_oauth_code(
+                'auth_code',
+                ['workspace-1'],
+                {'workspace-1': 1_700_000_000},
+            )
 
         # Verify
         assert result['access_token'] == 'new_access_token'
+        assert mock_session_obj.post.call_args.kwargs['json'] == {
+            'code': 'auth_code',
+            'instance_id': constants.instance_id,
+            'workspace_uuids': ['workspace-1'],
+            'workspace_created_ats': {'workspace-1': 1_700_000_000},
+        }
 
     async def test_exchange_oauth_code_api_error(self):
         """Raises ValueError on API error."""

tests/unit_tests/telemetry/test_heartbeat.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import json
+from datetime import datetime, timezone
 from types import SimpleNamespace
 
 import pytest
@@ -14,6 +15,12 @@ def get_heartbeat_module():
     return import_module('langbot.pkg.telemetry.heartbeat')
 
 
+def test_workspace_created_timestamp_treats_naive_database_values_as_utc():
+    heartbeat = get_heartbeat_module()
+    created_at = datetime(2026, 8, 4, 0, 0, 0)
+    assert heartbeat._workspace_created_timestamp(created_at) == 1785801600
+
+
 def make_app():
     ap = Mock()
     ap.instance_config = Mock()
@@ -57,15 +64,17 @@ def make_app():
 
 class TestBuildHeartbeatPayload:
     @pytest.mark.asyncio
-    async def test_payload_shape(self):
+    async def test_payload_shape(self, monkeypatch):
         heartbeat = get_heartbeat_module()
+        monkeypatch.setattr(heartbeat.constants, 'instance_id', 'instance-test')
         ap = make_app()
         payload = await heartbeat.build_heartbeat_payload(ap, workspace_uuid='workspace-a')
 
         assert payload['event_type'] == 'instance_heartbeat'
         assert payload['query_id'] == ''
         assert payload['workspace_uuid'] == 'workspace-a'
-        assert 'instance_id' not in payload
+        assert payload['instance_id']
+        assert payload['workspace_create_ts'] == 0
         assert 'instance_create_ts' in payload
         assert 'timestamp' in payload
         f = payload['features']
@@ -100,8 +109,9 @@ class TestBuildHeartbeatPayload:
         assert payload['features']['pipeline_count'] == -1
 
     @pytest.mark.asyncio
-    async def test_cloud_counts_loaded_registries_without_tenant_sql(self):
+    async def test_cloud_counts_loaded_registries_without_tenant_sql(self, monkeypatch):
         heartbeat = get_heartbeat_module()
+        monkeypatch.setattr(heartbeat.constants, 'instance_id', 'instance-test')
         ap = make_app()
         ap.persistence_mgr.mode = SimpleNamespace(value='cloud_runtime')
         ap.persistence_mgr.execute_async = AsyncMock(
@@ -140,7 +150,11 @@ class TestBuildHeartbeatPayload:
         ap.workspace_service.list_active_execution_bindings = AsyncMock(
             return_value=[
                 SimpleNamespace(workspace_uuid='workspace-a', placement_generation=7),
-                SimpleNamespace(workspace_uuid='workspace-b', placement_generation=9),
+                SimpleNamespace(
+                    workspace_uuid='workspace-b',
+                    placement_generation=9,
+                    workspace_created_at=datetime(2026, 8, 4, tzinfo=timezone.utc),
+                ),
             ],
         )
         ap.platform_mgr._bots_by_key[('instance-a', 'workspace-b', 'bot-b')] = SimpleNamespace(
@@ -150,7 +164,9 @@ class TestBuildHeartbeatPayload:
         payloads = await heartbeat.build_heartbeat_payloads(ap)
 
         assert [payload['workspace_uuid'] for payload in payloads] == ['workspace-a', 'workspace-b']
-        assert all('instance_id' not in payload for payload in payloads)
+        assert all(payload['instance_id'] for payload in payloads)
+        assert payloads[0]['workspace_create_ts'] == 0
+        assert payloads[1]['workspace_create_ts'] == 1785801600
         by_workspace = {payload['workspace_uuid']: payload['features'] for payload in payloads}
         assert by_workspace['workspace-a']['pipeline_count'] == 2
         assert by_workspace['workspace-a']['mcp_server_count'] == 3

tests/unit_tests/telemetry/test_telemetry.py

@@ -596,6 +596,36 @@ class TestTelemetryManagedRuntimeAuthentication:
         assert captured['headers'] == {'X-LangBot-Telemetry-Token': 'managed-runtime-secret'}
 
 
+class TestAuthenticatedWorkspaceReporter:
+    @pytest.mark.asyncio
+    async def test_workspace_owner_access_token_is_sent_as_bearer(self):
+        telemetry = get_telemetry_module()
+        mock_app = Mock()
+        mock_app.logger = Mock()
+        mock_app.user_service = Mock()
+        mock_app.user_service.get_workspace_owner = AsyncMock(
+            return_value=Mock(user='owner@example.com', space_access_token='expired-token')
+        )
+        mock_app.space_service = Mock()
+        mock_app.space_service.get_valid_access_token = AsyncMock(return_value='refreshed-workspace-owner-token')
+        manager = telemetry.TelemetryManager(mock_app)
+        manager.telemetry_config = {'url': 'https://example.com'}
+
+        response = Mock(status_code=200, text='')
+        response.json = Mock(return_value={'code': 0})
+        mock_client = Mock()
+        mock_client.post = Mock(return_value=response)
+
+        with patch.object(httpx, 'AsyncClient', return_value=mock_client):
+            await manager.send({'query_id': 'q-1', 'workspace_uuid': 'workspace-1'})
+
+        mock_app.user_service.get_workspace_owner.assert_awaited_once_with('workspace-1')
+        mock_app.space_service.get_valid_access_token.assert_awaited_once_with('owner@example.com')
+        assert mock_client.post.call_args.kwargs['headers'] == {
+            'Authorization': 'Bearer refreshed-workspace-owner-token'
+        }
+
+
 class TestStartSendTask:
     """Tests for start_send_task() method."""
 

tests/unit_tests/telemetry/test_workspace_identity.py

@@ -7,25 +7,28 @@ from types import SimpleNamespace
 def test_standard_oss_instance_id_aligns_to_embedded_uuid():
     from langbot.pkg.workspace.identity import workspace_uuid_from_instance_id
 
-    instance_uuid = "a711d9e4-0953-443f-a0e9-7dd50193a79f"
+    instance_uuid = 'a711d9e4-0953-443f-a0e9-7dd50193a79f'
 
     assert workspace_uuid_from_instance_id(instance_uuid) == instance_uuid
-    assert workspace_uuid_from_instance_id(f"instance_{instance_uuid}") == instance_uuid
+    assert workspace_uuid_from_instance_id(f'instance_{instance_uuid}') == instance_uuid
 
 
 def test_custom_legacy_instance_id_maps_to_stable_valid_uuid():
     from langbot.pkg.workspace.identity import workspace_uuid_from_instance_id
 
-    first = workspace_uuid_from_instance_id("instance_migration_test")
-    second = workspace_uuid_from_instance_id("instance_migration_test")
+    first = workspace_uuid_from_instance_id('instance_migration_test')
+    second = workspace_uuid_from_instance_id('instance_migration_test')
 
     assert first == second
     assert str(uuid.UUID(first)) == first
 
 
-def test_query_telemetry_identity_uses_execution_workspace_only():
+def test_query_telemetry_identity_reports_instance_and_workspace():
     from langbot.pkg.telemetry.identity import workspace_identity
 
-    identity = workspace_identity(SimpleNamespace(workspace_uuid="workspace-a", instance_uuid="instance-a"))
+    identity = workspace_identity(SimpleNamespace(workspace_uuid='workspace-a', instance_uuid='instance-a'))
 
-    assert identity == {"workspace_uuid": "workspace-a"}
+    assert identity == {
+        'instance_id': 'instance-a',
+        'workspace_uuid': 'workspace-a',
+    }