mainsequence 9.0.21.dev64__tar.gz → 9.0.22.dev65__tar.gz

{mainsequence-9.0.21.dev64 → mainsequence-9.0.22.dev65}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mainsequence
-Version: 9.0.21.dev64
+Version: 9.0.22.dev65
 Summary: Main Sequence SDK
 Author-email: Main Sequence GmbH <dev@main-sequence.io>
 License-Expression: MIT

{mainsequence-9.0.21.dev64 → mainsequence-9.0.22.dev65}/mainsequence/cli/api.py

@@ -76,7 +76,7 @@ class _CliSession(requests.Session):
         if kwargs.get("timeout") is not None:
             return super().request(method, url, **kwargs)
         # The CLI imports the SDK client on its first request, not when it starts.
-        from mainsequence.client._transport import DEFAULT_TIMEOUT, request_budget
+        from mainsequence._transport import DEFAULT_TIMEOUT, request_budget
 
         with request_budget(DEFAULT_TIMEOUT) as budget:
             kwargs["timeout"] = budget.timeout()
@@ -514,7 +514,7 @@ def _mcp_json_rpc(
     default request budget: 5 s to connect, 120 s to read, 125 s in total.
     """
     # The CLI imports the SDK client only where it needs it.
-    from mainsequence.client._transport import DEFAULT_TIMEOUT, request_budget
+    from mainsequence._transport import DEFAULT_TIMEOUT, request_budget
 
     payload = {
         "jsonrpc": "2.0",

{mainsequence-9.0.21.dev64 → mainsequence-9.0.22.dev65}/mainsequence/client/agent_runtime_models.py

@@ -12,7 +12,7 @@ from typing import Any, ClassVar, Literal
 import requests
 from pydantic import ConfigDict, Field, model_validator
 
-from ._transport import request_budget, request_within_budget
+from .._transport import request_budget, request_within_budget
 from .base import (
     BaseObjectOrm,
     BasePydanticModel,

{mainsequence-9.0.21.dev64 → mainsequence-9.0.22.dev65}/mainsequence/client/utils.py

@@ -28,10 +28,10 @@ from mainsequence.runtime_credential_exchange import (
     identity_token_file_from_environment,
 )
 
-from ._transport import DEFAULT_ALLOWED_METHODS as DEFAULT_ALLOWED_METHODS
-from ._transport import DEFAULT_STATUS_FORCELIST as DEFAULT_STATUS_FORCELIST
-from ._transport import DEFAULT_TIMEOUT as DEFAULT_TIMEOUT
-from ._transport import DeadlineHTTPAdapter, remaining_timeout, request_budget
+from .._transport import DEFAULT_ALLOWED_METHODS as DEFAULT_ALLOWED_METHODS
+from .._transport import DEFAULT_STATUS_FORCELIST as DEFAULT_STATUS_FORCELIST
+from .._transport import DEFAULT_TIMEOUT as DEFAULT_TIMEOUT
+from .._transport import DeadlineHTTPAdapter, remaining_timeout, request_budget
 
 # ---- Backend defaults (single source of truth) ----
 MAINSEQUENCE_ENDPOINT = resolve_backend_endpoint()

{mainsequence-9.0.21.dev64 → mainsequence-9.0.22.dev65}/mainsequence/logconf.py

@@ -5,7 +5,6 @@ import logging
 import logging.config
 import os
 import sys
-import time
 import traceback
 from collections.abc import Mapping
 from importlib.metadata import PackageNotFoundError
@@ -20,6 +19,7 @@ from structlog.contextvars import bind_contextvars, unbind_contextvars
 from structlog.dev import ConsoleRenderer
 from structlog.stdlib import BoundLogger
 
+from ._transport import DeadlineHTTPAdapter, request_budget
 from .defaults import resolve_backend_endpoint
 from .instrumentation import OTelJSONRenderer
 from .repository_identity_security import UNSUPPORTED_SOURCE_IDENTITY_ENV_NAMES
@@ -76,14 +76,27 @@ class CustomConsoleRenderer(ConsoleRenderer):
         return rendered
 
 
-def _request_job_startup_state(*, timeout_s: float = 10.0) -> dict[str, Any]:
+def _request_job_startup_state(*, timeout_s: float | None = None) -> dict[str, Any]:
     """
     Fetch startup state from backend using current env vars (JWT auth, endpoint, job_run_uid).
     Safe to call later after auth when access/refresh tokens become available.
+
+    The whole read, token renewal included, ends within one request budget:
+    ``timeout_s`` seconds in total, or the SDK's default request budget when it is
+    None. The GET is retried like any other SDK read. The startup state only
+    enriches the logger, so when the backend cannot be reached or does not answer
+    in time this reports it on stderr and returns an empty state instead of
+    failing the import.
     """
     if not is_running_in_pod():
         return {}
 
+    with request_budget(timeout_s) as budget:
+        return _read_job_startup_state(budget)
+
+
+def _read_job_startup_state(budget) -> dict[str, Any]:
+
     auth_mode = (os.getenv("MAINSEQUENCE_AUTH_MODE") or "jwt").strip().lower()
 
     def _backend_base_url() -> str:
@@ -110,7 +123,7 @@ def _request_job_startup_state(*, timeout_s: float = 10.0) -> dict[str, Any]:
                 token_url,
                 headers={"Content-Type": "application/json"},
                 json=body,
-                timeout=timeout_s,
+                timeout=budget.timeout(),
                 allow_redirects=False,
             )
 
@@ -119,7 +132,7 @@ def _request_job_startup_state(*, timeout_s: float = 10.0) -> dict[str, Any]:
                 _post,
                 credential_id=(os.getenv(RUNTIME_CREDENTIAL_ID_ENV) or "").strip(),
                 identity_token_file=identity_token_file_from_environment(),
-                deadline=time.monotonic() + timeout_s,
+                deadline=budget.deadline,
             )
         except Exception:
             return False
@@ -155,7 +168,7 @@ def _request_job_startup_state(*, timeout_s: float = 10.0) -> dict[str, Any]:
                 f"{_backend_base_url()}/auth/jwt-token/token/refresh/",
                 headers={"Content-Type": "application/json"},
                 json={"refresh": refresh_token},
-                timeout=timeout_s,
+                timeout=budget.timeout(),
             )
         except Exception:
             return False
@@ -202,10 +215,25 @@ def _request_job_startup_state(*, timeout_s: float = 10.0) -> dict[str, Any]:
     endpoint = f"{_backend_base_url()}/api/v1/job-runs/{job_run_uid}/startup-state/"
     params: dict[str, Any] = {}
 
-    resp = requests.get(endpoint, headers=headers, params=params, timeout=timeout_s)
-    if resp.status_code in {401, 403} and using_jwt and _refresh_access_token():
-        headers, _ = _auth_headers()
-        resp = requests.get(endpoint, headers=headers, params=params, timeout=timeout_s)
+    try:
+        with requests.Session() as session:
+            adapter = DeadlineHTTPAdapter(retries=3, backoff_factor=0.5)
+            session.mount("https://", adapter)
+            session.mount("http://", adapter)
+            resp = session.get(endpoint, headers=headers, params=params, timeout=budget.timeout())
+            if resp.status_code in {401, 403} and using_jwt and _refresh_access_token():
+                headers, _ = _auth_headers()
+                resp = session.get(
+                    endpoint, headers=headers, params=params, timeout=budget.timeout()
+                )
+    except requests.RequestException as exc:
+        # The exception text is not echoed: it can carry the request URL and headers.
+        print(
+            "mainsequence: the Job startup state could not be read "
+            f"({type(exc).__name__}); continuing without it.",
+            file=sys.stderr,
+        )
+        return {}
 
     if resp.status_code != 200:
         # Do not spam stdout during CLI usage when auth context is not present yet.
@@ -426,7 +454,7 @@ def build_application_logger(application_name: str = "ms-sdk", **metadata):
     return logger
 
 
-def refresh_application_logger_bindings(*, timeout_s: float = 10.0) -> BoundLogger:
+def refresh_application_logger_bindings(*, timeout_s: float | None = None) -> BoundLogger:
     """
     Importable helper: call this AFTER the user authenticates / token exists.
     It re-fetches backend state and binds the same keys as build_application_logger().

{mainsequence-9.0.21.dev64 → mainsequence-9.0.22.dev65}/mainsequence.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mainsequence
-Version: 9.0.21.dev64
+Version: 9.0.22.dev65
 Summary: Main Sequence SDK
 Author-email: Main Sequence GmbH <dev@main-sequence.io>
 License-Expression: MIT

{mainsequence-9.0.21.dev64 → mainsequence-9.0.22.dev65}/mainsequence.egg-info/SOURCES.txt

@@ -14,6 +14,7 @@ agent_scaffold/skills/sdk_code_repository_execution/SKILL.md
 mainsequence/__init__.py
 mainsequence/__main__.py
 mainsequence/_request_identity.py
+mainsequence/_transport.py
 mainsequence/bootstrap.py
 mainsequence/code_repository_context.py
 mainsequence/code_repository_skills.py
@@ -44,7 +45,6 @@ mainsequence/cli/sdk_utils.py
 mainsequence/cli/ssh_utils.py
 mainsequence/cli/ui.py
 mainsequence/client/__init__.py
-mainsequence/client/_transport.py
 mainsequence/client/agent_runtime_models.py
 mainsequence/client/base.py
 mainsequence/client/client.py

{mainsequence-9.0.21.dev64 → mainsequence-9.0.22.dev65}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "mainsequence"
-version = "9.0.21.dev64"
+version = "9.0.22.dev65"
 description = "Main Sequence SDK"
 readme = "README.md"
 requires-python = ">=3.13"

{mainsequence-9.0.21.dev64 → mainsequence-9.0.22.dev65}/tests/test_auth_precedence.py

@@ -22,6 +22,12 @@ _UNSUPPORTED_ENVIRONMENT_UID_ENV = (
 )
 
 
+def _stub_startup_state_get(monkeypatch, fake_get):
+    monkeypatch.setattr(
+        requests.Session, "get", lambda _session, url, **kwargs: fake_get(url, **kwargs)
+    )
+
+
 class _FakeResponse:
     def __init__(self, status_code: int, payload: dict | None = None):
         self.status_code = status_code
@@ -107,7 +113,7 @@ def test_logconf_refreshes_jwt_before_startup_state_request(monkeypatch):
         post_calls.append({"url": url, **kwargs})
         return _FakeResponse(200, {"access": "jwt-access"})
 
-    monkeypatch.setattr(requests, "get", _fake_get)
+    _stub_startup_state_get(monkeypatch, _fake_get)
     monkeypatch.setattr(requests, "post", _fake_post)
 
     _load_mainsequence_submodule("mainsequence.logconf")
@@ -143,11 +149,14 @@ def test_logconf_runtime_credential_exchanges_before_startup_state_request(monke
             },
         )
 
-    monkeypatch.setattr(requests, "get", _fake_get)
+    _stub_startup_state_get(monkeypatch, _fake_get)
     monkeypatch.setattr(requests, "post", _fake_post)
 
     _load_mainsequence_submodule("mainsequence.logconf")
 
+    # The exchange receives what is left of the startup-state read's request budget.
+    timeout = post_calls[0].pop("timeout")
+    assert 0 < timeout.total <= 125.0
     assert post_calls == [
         {
             "url": "https://backend.example/api/v1/runtime-credentials/token/",
@@ -156,7 +165,6 @@ def test_logconf_runtime_credential_exchanges_before_startup_state_request(monke
                 "credential_id": "cred-id",
                 "workload_identity_token": "projected-token",
             },
-            "timeout": 10.0,
             "allow_redirects": False,
         }
     ]
@@ -192,7 +200,7 @@ def test_logconf_runtime_credential_retries_after_auth_failure(monkeypatch, tmp_
             },
         )
 
-    monkeypatch.setattr(requests, "get", _fake_get)
+    _stub_startup_state_get(monkeypatch, _fake_get)
     monkeypatch.setattr(requests, "post", _fake_post)
 
     _load_mainsequence_submodule("mainsequence.logconf")
@@ -594,7 +602,7 @@ def test_logconf_session_jwt_does_not_refresh(monkeypatch):
         post_calls.append({"url": url, **kwargs})
         return _FakeResponse(200, {"access": "unexpected"})
 
-    monkeypatch.setattr(requests, "get", _fake_get)
+    _stub_startup_state_get(monkeypatch, _fake_get)
     monkeypatch.setattr(requests, "post", _fake_post)
 
     _load_mainsequence_submodule("mainsequence.logconf")

{mainsequence-9.0.21.dev64 → mainsequence-9.0.22.dev65}/tests/test_cli_request_budget.py

@@ -15,9 +15,9 @@ import pytest
 import requests
 from typer.testing import CliRunner
 
+from mainsequence import _transport as transport
 from mainsequence.cli import api as cli_api
 from mainsequence.cli import cli as cli_mod
-from mainsequence.client import _transport as transport
 from tests._http_server import answer_json, answer_sequence, never_answer, serve
 
 runner = CliRunner()

{mainsequence-9.0.21.dev64 → mainsequence-9.0.22.dev65}/tests/test_code_repository_skills.py

@@ -8,9 +8,9 @@ from pathlib import Path
 
 import pytest
 
+from mainsequence import _transport as transport
 from mainsequence.cli import api as cli_api
 from mainsequence.cli.api import ApiError
-from mainsequence.client import _transport as transport
 from mainsequence.code_repository_skills import (
     PLATFORM_ONTOLOGY_URI,
     PLATFORM_SKILL_URI_PREFIX,

{mainsequence-9.0.21.dev64 → mainsequence-9.0.22.dev65}/tests/test_http_retry_budget.py

@@ -7,7 +7,7 @@ from unittest.mock import Mock, PropertyMock
 import pytest
 import requests
 
-from mainsequence.client import _transport as transport
+from mainsequence import _transport as transport
 from mainsequence.client import utils
 from mainsequence.client.client import MainSequenceClient, MainSequenceClientConfig
 from mainsequence.client.exceptions import (

{mainsequence-9.0.21.dev64 → mainsequence-9.0.22.dev65}/tests/test_logconf.py

@@ -1,10 +1,14 @@
 from __future__ import annotations
 
 import os
+import sys
+import time
 
 import pytest
 import requests
+from urllib3.util import Timeout
 
+from tests._http_server import answer_after, answer_json, answer_sequence, never_answer, serve
 from tests._support import RETIRED_RUNTIME_CONTEXT_ENV_NAMES as _RUNTIME_CONTEXT_ENV_NAMES
 from tests._support import load_sdk_submodule as _load_mainsequence_submodule
 
@@ -40,9 +44,9 @@ def test_logconf_import_skips_job_startup_state_request_outside_pod(monkeypatch)
 
     def _fake_get(*args, **kwargs):
         calls.append((args, kwargs))
-        raise AssertionError("requests.get should not be called outside pod runtime")
+        raise AssertionError("the startup state should not be read outside pod runtime")
 
-    monkeypatch.setattr(requests, "get", _fake_get)
+    monkeypatch.setattr(requests.Session, "get", lambda _s, *a, **kw: _fake_get(*a, **kw))
 
     logconf = _load_mainsequence_submodule("mainsequence.logconf")
 
@@ -92,9 +96,9 @@ def test_logconf_import_skips_job_startup_state_request_without_job_run_uid(monk
 
     def _fake_get(*args, **kwargs):
         calls.append((args, kwargs))
-        raise AssertionError("requests.get should not be called without JOB_RUN_UID")
+        raise AssertionError("the startup state should not be read without JOB_RUN_UID")
 
-    monkeypatch.setattr(requests, "get", _fake_get)
+    monkeypatch.setattr(requests.Session, "get", lambda _s, *a, **kw: _fake_get(*a, **kw))
 
     logconf = _load_mainsequence_submodule("mainsequence.logconf")
 
@@ -132,7 +136,7 @@ def test_logconf_import_requests_job_run_detail_startup_state(monkeypatch):
         )
         return _FakeResponse()
 
-    monkeypatch.setattr(requests, "get", _fake_get)
+    monkeypatch.setattr(requests.Session, "get", lambda _s, *a, **kw: _fake_get(*a, **kw))
 
     logconf = _load_mainsequence_submodule("mainsequence.logconf")
 
@@ -147,3 +151,114 @@ def test_logconf_import_requests_job_run_detail_startup_state(monkeypatch):
     assert bindings["job_run_uid"] == "aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa"
     assert "project_id" not in bindings
     assert "data_source_id" not in bindings
+
+
+# The Job startup-state read (#215) crosses the real HTTP boundary with a local
+# server. A short stand-in for the SDK's default request budget keeps it fast.
+
+JOB_RUN_UID = "aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa"
+STARTUP_STATE_PATH = f"/api/v1/job-runs/{JOB_RUN_UID}/startup-state/"
+REFRESH_PATH = "/auth/jwt-token/token/refresh/"
+
+
+@pytest.fixture
+def job_logconf(monkeypatch):
+    """Load logconf outside a Job, then return it with the Job's environment set."""
+    for name in ("MAINSEQUENCE_ACCESS_TOKEN", "MAINSEQUENCE_REFRESH_TOKEN", "JOB_RUN_UID"):
+        monkeypatch.delenv(name, raising=False)
+    logconf = _load_mainsequence_submodule("mainsequence.logconf")
+    transport = sys.modules["mainsequence._transport"]
+
+    def start(base_url, *, auth_mode="session_jwt", default_timeout=(0.2, 0.3)):
+        monkeypatch.setattr(transport, "DEFAULT_TIMEOUT", default_timeout)
+        monkeypatch.setenv("MAINSEQUENCE_ENDPOINT", base_url.rstrip("/"))
+        monkeypatch.setenv("MAINSEQUENCE_AUTH_MODE", auth_mode)
+        monkeypatch.setenv("MAINSEQUENCE_ACCESS_TOKEN", "secret-access-token")
+        monkeypatch.setenv("JOB_RUN_UID", JOB_RUN_UID)
+        return logconf
+
+    return start
+
+
+def _startup_gets(calls):
+    return [call for call in calls if call[:2] == ("GET", STARTUP_STATE_PATH)]
+
+
+def test_logger_builds_when_the_startup_state_read_times_out(job_logconf, capsys):
+    with serve({("GET", STARTUP_STATE_PATH): never_answer()}) as (url, calls):
+        logconf = job_logconf(url)
+        started = time.monotonic()
+        logger = logconf.build_application_logger()
+        elapsed = time.monotonic() - started
+
+    assert elapsed < 0.5 + 0.3
+    assert len(_startup_gets(calls)) == 1
+    bound = logconf.dump_structlog_bound_logger(logger)["bound_context"]
+    assert "job_run_uid" not in bound
+    err = capsys.readouterr().err
+    assert "the Job startup state could not be read (ReadTimeout)" in err
+    assert "secret-access-token" not in err
+
+
+def test_logger_builds_when_the_backend_cannot_be_reached(job_logconf, capsys):
+    with serve({}) as (url, _calls):
+        pass  # the port is closed again once the server stops
+    logconf = job_logconf(url)
+
+    logger = logconf.build_application_logger()
+
+    assert logger is not None
+    assert "could not be read (ConnectionError)" in capsys.readouterr().err
+
+
+def test_startup_state_read_retries_a_failed_safe_read_within_its_budget(job_logconf):
+    state = {"job_run_uid": JOB_RUN_UID, "command_id": 12}
+    route = answer_sequence(
+        answer_after(1.0, 200, {}), answer_json(503, {}), answer_json(200, state)
+    )
+    with serve({("GET", STARTUP_STATE_PATH): route}) as (url, calls):
+        logconf = job_logconf(url, default_timeout=Timeout(connect=0.3, read=0.3, total=3.0))
+        logger = logconf.build_application_logger()
+
+    assert len(_startup_gets(calls)) == 3
+    bound = logconf.dump_structlog_bound_logger(logger)["bound_context"]
+    assert bound["job_run_uid"] == JOB_RUN_UID
+    assert bound["command_id"] == 12
+
+
+def test_denied_startup_state_is_not_treated_as_granted(job_logconf, capsys):
+    route = answer_json(403, {"detail": "denied"})
+    with serve({("GET", STARTUP_STATE_PATH): route}) as (url, calls):
+        logconf = job_logconf(url)
+        assert logconf._request_job_startup_state() == {}
+
+    assert len(_startup_gets(calls)) == 1
+    assert capsys.readouterr().err == ""
+
+
+def test_failed_token_renewal_after_a_denial_shares_the_budget(job_logconf, monkeypatch):
+    routes = {
+        ("GET", STARTUP_STATE_PATH): answer_json(401, {}),
+        ("POST", REFRESH_PATH): never_answer(),
+    }
+    with serve(routes) as (url, calls):
+        logconf = job_logconf(url, auth_mode="jwt")
+        monkeypatch.setenv("MAINSEQUENCE_REFRESH_TOKEN", "secret-refresh-token")
+        started = time.monotonic()
+        assert logconf._request_job_startup_state() == {}
+        elapsed = time.monotonic() - started
+
+    assert elapsed < 0.5 + 0.3
+    assert [call[:2] for call in calls] == [("GET", STARTUP_STATE_PATH), ("POST", REFRESH_PATH)]
+
+
+def test_refreshing_logger_bindings_honours_an_explicit_total_timeout(job_logconf, capsys):
+    with serve({("GET", STARTUP_STATE_PATH): never_answer()}) as (url, _calls):
+        logconf = job_logconf(url, default_timeout=(5.0, 120.0))
+        started = time.monotonic()
+        logger = logconf.refresh_application_logger_bindings(timeout_s=0.3)
+        elapsed = time.monotonic() - started
+
+    assert logger is not None
+    assert elapsed < 0.3 + 0.3
+    assert "could not be read (ReadTimeout)" in capsys.readouterr().err

{mainsequence-9.0.21.dev64 → mainsequence-9.0.22.dev65}/tests/test_runtime_credential_exchange.py

@@ -138,7 +138,7 @@ def test_no_code_path_reads_the_secret_variable(monkeypatch, runtime):
     # os.getenv() reads os.environ, so every lookup by name goes through the recorder.
     monkeypatch.setattr(os, "environ", environ)
     exchange = _Exchange(monkeypatch, _granted(), _granted())
-    monkeypatch.setattr(requests, "get", lambda url, **kwargs: _Answer(200, {}))
+    monkeypatch.setattr(requests.Session, "get", lambda _session, url, **kwargs: _Answer(200, {}))
 
     # Importing the client configures the logger, whose start-up exchange runs first.
     utils = load_sdk_submodule("mainsequence.client.utils")
@@ -306,7 +306,7 @@ def test_exchange_retries_are_bounded(monkeypatch, runtime):
 def test_each_attempt_receives_only_the_remaining_budget(monkeypatch, runtime):
     _use_token_file(monkeypatch, runtime)
     utils = load_sdk_submodule("mainsequence.client.utils")
-    transport = importlib.import_module("mainsequence.client._transport")
+    transport = importlib.import_module("mainsequence._transport")
     exchange = _Exchange(monkeypatch, _Answer(503, headers={"Retry-After": "4"}), _granted())
     now = [100.0]
 
@@ -434,7 +434,7 @@ def test_logger_start_up_exchange_sends_the_identity_token_and_retries(monkeypat
         authorizations.append(kwargs["headers"].get("Authorization"))
         return _Answer(200, {"job_run_uid": JOB_RUN_UID})
 
-    monkeypatch.setattr(requests, "get", _get)
+    monkeypatch.setattr(requests.Session, "get", lambda _session, url, **kw: _get(url, **kw))
 
     load_sdk_submodule("mainsequence.logconf")
 
@@ -452,7 +452,7 @@ def test_logger_start_up_exchange_needs_the_identity_token_file(monkeypatch, run
     monkeypatch.setenv(SECRET_ENV, SECRET)
     monkeypatch.setenv("JOB_RUN_UID", JOB_RUN_UID)
     exchange = _Exchange(monkeypatch)
-    monkeypatch.setattr(requests, "get", lambda url, **kwargs: _Answer(401))
+    monkeypatch.setattr(requests.Session, "get", lambda _session, url, **kwargs: _Answer(401))
 
     load_sdk_submodule("mainsequence.logconf")