diff --git a/Claudette.sublime-settings b/Claudette.sublime-settings index 49e4e31..66abd70 100644 --- a/Claudette.sublime-settings +++ b/Claudette.sublime-settings @@ -1,4 +1,23 @@ { + // Provider: "anthropic" (default) or "bedrock" (AWS Bedrock). + // When set to "bedrock", requests are signed with AWS SigV4 and sent + // to the Bedrock runtime endpoint. No api_key is required. + "provider": "anthropic", + // AWS region for Bedrock (only used when provider is "bedrock"). + "aws_region": "us-east-1", + // AWS credentials for Bedrock. Resolution order: + // 1. aws_access_key_id + aws_secret_access_key (+ optional aws_session_token) + // 2. AWS_ACCESS_KEY_ID / AWS_SECRET_ACCESS_KEY environment variables + // 3. aws_profile setting (calls `aws configure export-credentials`) + // 4. Default profile via `aws configure export-credentials` + // Steps 3 and 4 require the AWS CLI v2 to be installed and on PATH. + // Resolved credentials are cached per profile until their reported + // Expiration (or for ~15 minutes when no expiration is reported). + // "aws_access_key_id": "", + // "aws_secret_access_key": "", + // "aws_session_token": "", + // "aws_profile": "", + // // API key configuration. You can define a single key or multiple keys. // // Example for a single key: diff --git a/api/api.py b/api/api.py index cd6132c..423f860 100644 --- a/api/api.py +++ b/api/api.py @@ -1,6 +1,5 @@ import json import os -import select import socket import ssl import urllib.error @@ -25,7 +24,13 @@ ) from ..utils import claudette_chat_status_message, claudette_get_api_key_value from . import session_stats -from .cancellation import CancellationToken +from .bedrock import ( + BEDROCK_ANTHROPIC_VERSION, + BedrockHTTPError, + bedrock_request, + get_aws_credentials, + parse_event_stream, +) from .errors import ( handle_model_not_found, is_model_not_found_error, @@ -39,6 +44,8 @@ parse_web_search_items, ) +KNOWN_PROVIDERS = ("anthropic", "bedrock") + class CancelledException(Exception): """Raised when a request is cancelled.""" @@ -49,8 +56,19 @@ class CancelledException(Exception): class ClaudetteClaudeAPI: def __init__(self): self.settings = sublime.load_settings(SETTINGS_FILE) + provider = self.settings.get("provider", "anthropic") + if provider not in KNOWN_PROVIDERS: + print( + "Claudette: unknown provider {0!r}; falling back to " + "'anthropic'. Valid values: {1}.".format( + provider, ", ".join(KNOWN_PROVIDERS) + ) + ) + provider = "anthropic" + self.provider = provider self.api_key = claudette_get_api_key_value() self.base_url = self.settings.get("base_url", DEFAULT_BASE_URL) + self.aws_region = self.settings.get("aws_region", "us-east-1") try: self.max_tokens = int(self.settings.get("max_tokens", MAX_TOKENS)) except (TypeError, ValueError): @@ -64,6 +82,22 @@ def __init__(self): self.pricing = self.settings.get("pricing") self.verify_ssl = self.settings.get("verify_ssl", DEFAULT_VERIFY_SSL) + @property + def _is_bedrock(self): + return self.provider == "bedrock" + + def _get_bedrock_credentials(self): + """Get AWS credentials, raising a clear error if unavailable.""" + credentials = get_aws_credentials(self.settings) + if not credentials: + raise RuntimeError( + "AWS credentials not found. Configure aws_profile, " + "aws_access_key_id/aws_secret_access_key in settings, " + "or set AWS_ACCESS_KEY_ID/AWS_SECRET_ACCESS_KEY " + "environment variables." + ) + return credentials + def _get_ssl_context(self): """Create and return an SSL context based on verify_ssl setting.""" if self.verify_ssl: @@ -213,24 +247,41 @@ def _request_non_streaming( reading the response body so the user isn't stuck waiting for the full HTTP timeout. """ - headers = { - "x-api-key": self.api_key, - "anthropic-version": ANTHROPIC_VERSION, - "content-type": "application/json", - } - headers.update(self._get_custom_headers()) - data = { "messages": messages, "max_tokens": self.max_tokens, - "model": self.model, - "stream": False, "system": system_messages, "temperature": self.get_valid_temperature(self.temperature), } if tools_list: data["tools"] = tools_list + if self._is_bedrock: + data["anthropic_version"] = BEDROCK_ANTHROPIC_VERSION + credentials = self._get_bedrock_credentials() + body = bedrock_request( + self.aws_region, self.model, data, credentials, + streaming=False, verify_ssl=self.verify_ssl, + ) + msg = ( + body.get("message") + if body.get("message") is not None + else body + ) + if not isinstance(msg, dict): + msg = {} + usage = body.get("usage") or msg.get("usage") or {} + return msg, usage + + data["model"] = self.model + data["stream"] = False + headers = { + "x-api-key": self.api_key, + "anthropic-version": ANTHROPIC_VERSION, + "content-type": "application/json", + } + headers.update(self._get_custom_headers()) + req = urllib.request.Request( urllib.parse.urljoin(self.base_url, "messages"), data=json.dumps(data).encode("utf-8"), @@ -304,7 +355,7 @@ def check_cancelled(): if not filtered: return - if not self.api_key: + if not self._is_bedrock and not self.api_key: handle_error( "[Error] The API key is not set. Please check your API key " "configuration." @@ -365,6 +416,21 @@ def check_cancelled(): return handle_error("[Error] {0}".format(error_message)) return + except BedrockHTTPError as e: + self.spinner.stop() + if chat_view_for_status: + sublime.set_timeout( + lambda: chat_view_for_status.clear_tool_status(), 0 + ) + if is_model_not_found_error( + e.status, e.error_type, e.message + ): + handle_model_not_found( + e.message, window, settings, handle_error + ) + return + handle_error("[Error] {0}".format(e.message)) + return except urllib.error.URLError as e: self.spinner.stop() if chat_view_for_status: @@ -638,7 +704,7 @@ def is_cancelled(): ): return - if not self.api_key: + if not self._is_bedrock and not self.api_key: handle_error( "[Error] The API key is not set. Please check your API key " "configuration." @@ -648,13 +714,6 @@ def is_cancelled(): try: self.spinner.start("Fetching response") - headers = { - "x-api-key": self.api_key, - "anthropic-version": ANTHROPIC_VERSION, - "content-type": "application/json", - } - headers.update(self._get_custom_headers()) - filtered_messages = [ msg for msg in messages if self._message_has_content(msg) ] @@ -664,8 +723,6 @@ def is_cancelled(): data = { "messages": filtered_messages, "max_tokens": self.max_tokens, - "model": self.model, - "stream": True, "system": system_messages, "temperature": self.get_valid_temperature(self.temperature), } @@ -674,12 +731,11 @@ def is_cancelled(): if web_search_tool: data["tools"] = [web_search_tool] - req = urllib.request.Request( - urllib.parse.urljoin(self.base_url, "messages"), - data=json.dumps(data).encode("utf-8"), - headers=headers, - method="POST", - ) + if self._is_bedrock: + data["anthropic_version"] = BEDROCK_ANTHROPIC_VERSION + else: + data["model"] = self.model + data["stream"] = True try: ssl_context = self._get_ssl_context() @@ -687,82 +743,75 @@ def is_cancelled(): stream_current_block_type = None stream_current_block_index = None - # Use a timeout so connection doesn't hang forever - with urllib.request.urlopen( - req, context=ssl_context, timeout=30 - ) as response: - # Get socket for select-based polling with timeout - # This allows us to periodically check for cancellation - sock = None + def _iter_events_sse(response): + """Yield parsed JSON dicts from Anthropic SSE stream. + + Sets a short socket timeout so blocking readline()s can + be interrupted; the cancellation token is polled between + lines and on each timeout retry. + """ try: - # Try to get the underlying socket - if hasattr(response.fp, "raw"): - raw = response.fp.raw - if hasattr(raw, "_sock"): - sock = raw._sock + response.fp._sock.settimeout(0.5) except Exception: pass - - # When we cannot extract the raw socket for select(), - # set a short read timeout so readline() doesn't block - # indefinitely — this lets us check cancellation often. - if sock is None: - try: - response.fp._sock.settimeout(0.5) - except Exception: - pass - while True: - # Check for cancellation before reading if is_cancelled(): - response.close() - sublime.set_timeout( - lambda: chunk_callback( - "", is_done=True, was_cancelled=True - ), - 0, - ) return - - # Use select to wait for data with timeout - if sock is not None: - try: - ready, _, _ = select.select( - [sock], [], [], 0.3 - ) - if not ready: - # Timeout, check cancellation and retry - continue - except (ValueError, OSError, TypeError): - # Socket issue, fall through to blocking read - sock = None - try: - response.fp._sock.settimeout(0.5) - except Exception: - pass - try: line = response.readline() except socket.timeout: - # Read timed out — loop back to check cancellation continue if not line: - # End of stream - break - + return if line.isspace(): continue - + chunk = line.decode("utf-8") + if not chunk.startswith("data: "): + continue + chunk = chunk[6:] + if chunk.strip() == "[DONE]": + return try: - chunk = line.decode("utf-8") - if not chunk.startswith("data: "): - continue + yield json.loads(chunk) + except (json.JSONDecodeError, ValueError): + continue - chunk = chunk[6:] # Remove 'data: ' prefix - if chunk.strip() == "[DONE]": - break + def _open_and_iter(): + """Open connection and return (response, event_iter).""" + if self._is_bedrock: + credentials = self._get_bedrock_credentials() + resp = bedrock_request( + self.aws_region, self.model, data, credentials, + streaming=True, verify_ssl=self.verify_ssl, + ) + return resp, parse_event_stream( + resp, should_cancel=is_cancelled + ) + else: + headers = { + "x-api-key": self.api_key, + "anthropic-version": ANTHROPIC_VERSION, + "content-type": "application/json", + } + headers.update(self._get_custom_headers()) + req = urllib.request.Request( + urllib.parse.urljoin(self.base_url, "messages"), + data=json.dumps(data).encode("utf-8"), + headers=headers, + method="POST", + ) + resp = urllib.request.urlopen( + req, context=ssl_context, timeout=30 + ) + return resp, _iter_events_sse(resp) - data = json.loads(chunk) + response, event_iter = _open_and_iter() + try: + for data in event_iter: + if is_cancelled(): + break + + try: # Get initial input tokens from message_start if data.get("type") == "message_start": @@ -1048,9 +1097,22 @@ def _send_citation( ) except Exception: - # Skip invalid chunks without error messages continue + # Loop exited. If it was due to cancellation, signal it + # so the chat view can clear the response heading and + # the active-request token. + if is_cancelled(): + sublime.set_timeout( + lambda: chunk_callback( + "", is_done=True, was_cancelled=True + ), + 0, + ) + finally: + if hasattr(response, "close"): + response.close() + except urllib.error.HTTPError as e: error_type, error_message = parse_api_error(e) if is_model_not_found_error(e.code, error_type, error_message): @@ -1061,16 +1123,54 @@ def _send_citation( else: handle_error("[Error] {0}".format(error_message)) except urllib.error.URLError as e: - handle_error(f"[Error] {str(e)}") + handle_error("[Error] {0}".format(str(e))) + except BedrockHTTPError as e: + if is_model_not_found_error( + e.status, e.error_type, e.message + ): + window = chat_view.window() if chat_view else None + handle_model_not_found( + e.message, window, self.settings, handle_error + ) + else: + handle_error("[Error] {0}".format(e.message)) + except RuntimeError as e: + handle_error("[Error] {0}".format(str(e))) finally: self.spinner.stop() except Exception as e: - handle_error(f"[Error] {str(e)}") + handle_error("[Error] {0}".format(str(e))) self.spinner.stop() def fetch_models(self): + if self._is_bedrock: + # Verified currently-shipping Bedrock model IDs (as of 2026-06). + # If the model you want isn't here (e.g. Opus 4.7/4.8 reachable + # only via InvokeModel without an ARN-versioned id), set the + # `model` setting directly. + return [ + "anthropic.claude-opus-4-6-v1", + "anthropic.claude-opus-4-5-20251101-v1:0", + "anthropic.claude-sonnet-4-6", + "anthropic.claude-sonnet-4-5-20250929-v1:0", + "anthropic.claude-haiku-4-5-20251001-v1:0", + "us.anthropic.claude-opus-4-6-v1", + "us.anthropic.claude-opus-4-5-20251101-v1:0", + "us.anthropic.claude-sonnet-4-6", + "us.anthropic.claude-sonnet-4-5-20250929-v1:0", + "us.anthropic.claude-haiku-4-5-20251001-v1:0", + "eu.anthropic.claude-opus-4-6-v1", + "eu.anthropic.claude-sonnet-4-6", + "eu.anthropic.claude-sonnet-4-5-20250929-v1:0", + "eu.anthropic.claude-haiku-4-5-20251001-v1:0", + "apac.anthropic.claude-opus-4-6-v1", + "global.anthropic.claude-opus-4-6-v1", + "global.anthropic.claude-sonnet-4-6", + "global.anthropic.claude-sonnet-4-5-20250929-v1:0", + ] + if not self.api_key: sublime.error_message( "The API key is undefined. Please check your API key " diff --git a/api/bedrock.py b/api/bedrock.py new file mode 100644 index 0000000..6b0c762 --- /dev/null +++ b/api/bedrock.py @@ -0,0 +1,592 @@ +"""AWS Bedrock provider: SigV4 signing, request, and event-stream parsing.""" + +import base64 +import hashlib +import hmac +import http.client +import json +import os +import socket +import ssl +import struct +import subprocess +import sys +import threading +import urllib.parse +from datetime import datetime, timedelta, timezone + +BEDROCK_ANTHROPIC_VERSION = "bedrock-2023-05-31" + +# Cache resolved AWS credentials per profile to avoid re-spawning the AWS CLI +# on every request. SSO/profile lookups can take 0.5–2s. +_CREDENTIALS_CACHE = {} +_CREDENTIALS_CACHE_LOCK = threading.Lock() +# How long to trust env/settings credentials before re-resolving (no expiry +# is reported for those, so we expire defensively). +_STATIC_CREDENTIALS_TTL = timedelta(minutes=15) + + +class BedrockHTTPError(Exception): + """Raised when AWS Bedrock returns a non-2xx HTTP response. + + Attributes: + status: The HTTP status code. + message: The parsed error message from the response body. + error_type: The Bedrock error type (e.g. "ValidationException"), if + present in the response body. + """ + + def __init__(self, status, message, error_type=""): + super().__init__(f"Bedrock HTTP {status}: {message}") + self.status = status + self.message = message + self.error_type = error_type + + +def _hmac_sha256(key, msg): + if isinstance(msg, str): + msg = msg.encode("utf-8") + if isinstance(key, str): + key = key.encode("utf-8") + return hmac.new(key, msg, hashlib.sha256).digest() + + +def _get_signature_key(secret_key, date_stamp, region, service): + k_date = _hmac_sha256(("AWS4" + secret_key).encode("utf-8"), date_stamp) + k_region = _hmac_sha256(k_date, region) + k_service = _hmac_sha256(k_region, service) + k_signing = _hmac_sha256(k_service, "aws4_request") + return k_signing + + +def _uri_encode_path_for_signing(path): + """URI-encode each path segment for the SigV4 canonical request. + + Per the SigV4 spec for non-S3 services, the canonical URI is the path + URI-encoded once, with the segment separator '/' preserved. The path + we receive already has the model id encoded by ``_build_request_path``; + the colons in Bedrock model ids (e.g. "...:0") survive that initial + encode as literal '%3A' here, which is what Bedrock expects. + """ + segments = path.split("/") + return "/".join(urllib.parse.quote(seg, safe="") for seg in segments) + + +def _strip_export_prefix(line): + """Remove a leading 'export ' from a shell variable assignment line.""" + return line[7:] if line.startswith("export ") else line + + +def _strip_surrounding_quotes(value): + """Strip a single matching pair of ' or " from value, if present.""" + if len(value) >= 2 and value[0] == value[-1] and value[0] in ("'", '"'): + return value[1:-1] + return value + + +def _subprocess_kwargs_no_window(timeout=10): + """Return subprocess.run kwargs that don't flash a console on Windows.""" + kwargs = {"capture_output": True, "text": True, "timeout": timeout} + if sys.platform == "win32": + # CREATE_NO_WINDOW = 0x08000000. Avoids a console flash when called + # from a GUI process like Sublime Text. + kwargs["creationflags"] = 0x08000000 + return kwargs + + +def _parse_iso8601_expiration(value): + """Parse an ISO 8601 expiration string into a timezone-aware datetime.""" + if not value: + return None + # AWS returns e.g. "2026-06-12T15:00:00+00:00" with --format json. The + # 'Z' suffix is also possible; normalise both. + value = value.replace("Z", "+00:00") + try: + return datetime.fromisoformat(value) + except (TypeError, ValueError): + return None + + +def _get_credentials_from_profile(profile=None): + """Resolve credentials via the AWS CLI's export-credentials command. + + Uses ``aws configure export-credentials --format process``. + + Args: + profile: Optional named profile. If None, the default profile is used. + + Returns: + A dict with keys access_key, secret_key, session_token (str or None), + and expiration (datetime or None), or None if the AWS CLI is missing + or the profile lookup failed. + """ + cmd = ["aws", "configure", "export-credentials", "--format", "process"] + if profile: + cmd.extend(["--profile", profile]) + try: + result = subprocess.run(cmd, **_subprocess_kwargs_no_window()) + except (subprocess.TimeoutExpired, FileNotFoundError, OSError): + return None + if result.returncode != 0: + return None + try: + data = json.loads(result.stdout) + except (json.JSONDecodeError, ValueError): + return _get_credentials_from_profile_env(profile) + access_key = data.get("AccessKeyId", "") + secret_key = data.get("SecretAccessKey", "") + session_token = data.get("SessionToken") or None + expiration = _parse_iso8601_expiration(data.get("Expiration")) + if access_key and secret_key: + return { + "access_key": access_key, + "secret_key": secret_key, + "session_token": session_token, + "expiration": expiration, + } + return None + + +def _get_credentials_from_profile_env(profile=None): + """Fallback: parse 'aws configure export-credentials --format env'.""" + cmd = ["aws", "configure", "export-credentials", "--format", "env"] + if profile: + cmd.extend(["--profile", profile]) + try: + result = subprocess.run(cmd, **_subprocess_kwargs_no_window()) + except (subprocess.TimeoutExpired, FileNotFoundError, OSError): + return None + if result.returncode != 0: + return None + creds = {} + for raw in result.stdout.strip().split("\n"): + if "=" not in raw: + continue + line = _strip_export_prefix(raw.strip()) + key, _, value = line.partition("=") + creds[key.strip()] = _strip_surrounding_quotes(value.strip()) + access_key = creds.get("AWS_ACCESS_KEY_ID", "") + secret_key = creds.get("AWS_SECRET_ACCESS_KEY", "") + session_token = creds.get("AWS_SESSION_TOKEN") or None + if access_key and secret_key: + return { + "access_key": access_key, + "secret_key": secret_key, + "session_token": session_token, + "expiration": None, + } + return None + + +def _get_credentials_from_env(): + """Return AWS credentials from environment variables, or None.""" + access_key = os.environ.get("AWS_ACCESS_KEY_ID", "") + secret_key = os.environ.get("AWS_SECRET_ACCESS_KEY", "") + session_token = os.environ.get("AWS_SESSION_TOKEN") or None + if access_key and secret_key: + return { + "access_key": access_key, + "secret_key": secret_key, + "session_token": session_token, + "expiration": None, + } + return None + + +def _credentials_still_valid(creds): + """Return True if cached credentials have not yet expired.""" + if not creds: + return False + expiration = creds.get("expiration") + if expiration is None: + # Static creds (env/settings) — trust for a bounded window. + cached_at = creds.get("cached_at") + if cached_at is None: + return False + return datetime.now(timezone.utc) - cached_at < _STATIC_CREDENTIALS_TTL + # Refresh a minute early to avoid races against AWS-side expiry. + return datetime.now(timezone.utc) + timedelta(seconds=60) < expiration + + +def get_aws_credentials(settings): + """Resolve AWS credentials from settings, env vars, or AWS CLI profile. + + Resolution order: + 1. aws_access_key_id + aws_secret_access_key in settings + 2. AWS_ACCESS_KEY_ID / AWS_SECRET_ACCESS_KEY environment variables + 3. ``aws_profile`` setting (via 'aws configure export-credentials') + 4. Default profile (same call, no --profile) + + Resolved credentials are cached per cache-key (settings vs env vs profile + name) until their reported expiration, or for 15 minutes when no + expiration is reported. + + Args: + settings: The plugin settings object (Sublime ``Settings``). + + Returns: + A dict with keys access_key, secret_key, session_token, and + expiration. Or ``None`` if no credentials could be resolved. + """ + settings_access = settings.get("aws_access_key_id", "") + settings_secret = settings.get("aws_secret_access_key", "") + if settings_access and settings_secret: + cache_key = ("settings", settings_access) + with _CREDENTIALS_CACHE_LOCK: + cached = _CREDENTIALS_CACHE.get(cache_key) + if _credentials_still_valid(cached): + return cached + session_token = settings.get("aws_session_token") or None + creds = { + "access_key": settings_access, + "secret_key": settings_secret, + "session_token": session_token, + "expiration": None, + "cached_at": datetime.now(timezone.utc), + } + _CREDENTIALS_CACHE[cache_key] = creds + return creds + + env_access = os.environ.get("AWS_ACCESS_KEY_ID", "") + if env_access: + cache_key = ("env", env_access) + with _CREDENTIALS_CACHE_LOCK: + cached = _CREDENTIALS_CACHE.get(cache_key) + if _credentials_still_valid(cached): + return cached + creds = _get_credentials_from_env() + if creds: + creds["cached_at"] = datetime.now(timezone.utc) + _CREDENTIALS_CACHE[cache_key] = creds + return creds + + profile = settings.get("aws_profile") or None + cache_key = ("profile", profile or "__default__") + with _CREDENTIALS_CACHE_LOCK: + cached = _CREDENTIALS_CACHE.get(cache_key) + if _credentials_still_valid(cached): + return cached + creds = _get_credentials_from_profile(profile) + if creds: + if creds.get("expiration") is None: + creds["cached_at"] = datetime.now(timezone.utc) + _CREDENTIALS_CACHE[cache_key] = creds + return creds + return None + + +def _build_request_path(model_id, streaming=False): + """Build the URL path for the Bedrock invoke endpoint.""" + endpoint = "invoke-with-response-stream" if streaming else "invoke" + encoded_model = urllib.parse.quote(model_id, safe="") + return f"/model/{encoded_model}/{endpoint}" + + +def _parse_bedrock_error_body(body_text): + """Extract (message, error_type) from a Bedrock error response body.""" + try: + data = json.loads(body_text) + except (json.JSONDecodeError, ValueError): + return body_text, "" + message = data.get("message") or data.get("Message") or body_text + error_type = data.get("__type") or data.get("type") or "" + # __type often looks like "com.amazon.coral...#ValidationException"; keep + # only the last segment for readability. + if "#" in error_type: + error_type = error_type.split("#", 1)[1] + return message, error_type + + +class _BedrockStreamResponse: + """Tiny wrapper that owns both the http.client response and connection. + + The streaming caller only needs ``read`` (for the parser) and ``close`` + (for cleanup). Bundling the connection here means callers don't have to + poke at private attributes to release resources. + """ + + def __init__(self, response, conn): + self._response = response + self._conn = conn + + def read(self, n=-1): + return self._response.read(n) + + @property + def fp(self): + return self._response.fp + + def close(self): + try: + self._response.close() + finally: + self._conn.close() + + +def bedrock_request( + region, + model_id, + body_dict, + credentials, + streaming=False, + verify_ssl=True, + timeout=30, +): + """Make a SigV4-signed request to AWS Bedrock. + + Args: + region: AWS region, e.g. "us-east-1". + model_id: Bedrock model id, e.g. + "anthropic.claude-3-5-sonnet-20241022-v2:0". + body_dict: The request body as a JSON-serialisable dict. + credentials: A dict with access_key, secret_key, and optional + session_token (as returned by ``get_aws_credentials``). + streaming: When True, hits the invoke-with-response-stream endpoint + and returns a stream wrapper. When False, hits invoke and + returns the parsed JSON response. + verify_ssl: When False, disables certificate verification. + timeout: Per-operation socket timeout in seconds. + + Returns: + For non-streaming: the parsed JSON response (dict). + For streaming: a ``_BedrockStreamResponse`` exposing ``read``/ + ``close``/``fp`` — pass to ``parse_event_stream``. + + Raises: + BedrockHTTPError: On non-2xx responses. + """ + host = f"bedrock-runtime.{region}.amazonaws.com" + request_path = _build_request_path(model_id, streaming=streaming) + body = json.dumps(body_dict).encode("utf-8") + + now = datetime.now(timezone.utc) + amz_date = now.strftime("%Y%m%dT%H%M%SZ") + date_stamp = now.strftime("%Y%m%d") + payload_hash = hashlib.sha256(body).hexdigest() + + signed_headers_dict = { + "content-type": "application/json", + "host": host, + "x-amz-content-sha256": payload_hash, + "x-amz-date": amz_date, + } + if credentials.get("session_token"): + signed_headers_dict["x-amz-security-token"] = ( + credentials["session_token"] + ) + + signed_header_keys = sorted(signed_headers_dict.keys()) + signed_headers_str = ";".join(signed_header_keys) + canonical_headers = "".join( + f"{k}:{signed_headers_dict[k]}\n" for k in signed_header_keys + ) + + canonical_uri = _uri_encode_path_for_signing(request_path) + service = "bedrock" + + canonical_request = "\n".join([ + "POST", + canonical_uri, + "", # empty query string + canonical_headers, + signed_headers_str, + payload_hash, + ]) + + credential_scope = f"{date_stamp}/{region}/{service}/aws4_request" + string_to_sign = "\n".join([ + "AWS4-HMAC-SHA256", + amz_date, + credential_scope, + hashlib.sha256(canonical_request.encode("utf-8")).hexdigest(), + ]) + + signing_key = _get_signature_key( + credentials["secret_key"], date_stamp, region, service + ) + signature = hmac.new( + signing_key, string_to_sign.encode("utf-8"), hashlib.sha256 + ).hexdigest() + + authorization = ( + f"AWS4-HMAC-SHA256 Credential={credentials['access_key']}/" + f"{credential_scope}, SignedHeaders={signed_headers_str}, " + f"Signature={signature}" + ) + + headers = { + "Content-Type": "application/json", + "Host": host, + "X-Amz-Date": amz_date, + "X-Amz-Content-Sha256": payload_hash, + "Authorization": authorization, + } + if credentials.get("session_token"): + headers["X-Amz-Security-Token"] = credentials["session_token"] + + if verify_ssl: + context = ssl.create_default_context() + else: + context = ssl.create_default_context() + context.check_hostname = False + context.verify_mode = ssl.CERT_NONE + + # Use http.client to avoid urllib's URL re-encoding of the colon in + # Bedrock model ids. + conn = http.client.HTTPSConnection(host, context=context, timeout=timeout) + try: + conn.request("POST", request_path, body=body, headers=headers) + response = conn.getresponse() + except Exception: + conn.close() + raise + + if response.status != 200: + error_body = response.read().decode("utf-8", errors="replace") + conn.close() + message, error_type = _parse_bedrock_error_body(error_body) + raise BedrockHTTPError(response.status, message, error_type) + + if not streaming: + try: + raw = response.read().decode("utf-8") + return json.loads(raw) + finally: + conn.close() + + return _BedrockStreamResponse(response, conn) + + +def parse_event_stream(response, should_cancel=None): + """Parse the AWS event-stream binary format from a streaming response. + + Yields parsed JSON dicts that match the Anthropic SSE event shape, so + callers can reuse the same handling code for both providers. + + The prelude CRC and trailing message CRC are intentionally not verified + — this is a small, from-scratch parser and the underlying TLS connection + already provides integrity. If you ever need stronger guarantees, switch + to botocore's eventstream parser. + + Args: + response: A response object exposing ``read(n)`` and (optionally) + ``fp`` for socket access; typically the + ``_BedrockStreamResponse`` returned by ``bedrock_request``. + should_cancel: Optional zero-arg callable. When provided, the parser + polls it between events and after partial reads, and returns + early if it returns truthy. The underlying socket is set to a + short timeout so blocking reads can be interrupted. + + Raises: + BedrockHTTPError: When the stream contains an exception event. + RuntimeError: If the stream ends mid-message. + """ + buf = b"" + + # Make blocking reads interruptible by cancellation polling. + if should_cancel is not None: + sock = getattr(getattr(response, "fp", None), "_sock", None) + if sock is not None: + try: + sock.settimeout(0.5) + except Exception: + pass + + def cancelled(): + return should_cancel is not None and should_cancel() + + def read_exactly(n): + nonlocal buf + while len(buf) < n: + if cancelled(): + return None + try: + chunk = response.read(n - len(buf)) + except socket.timeout: + continue + if not chunk: + if buf: + raise RuntimeError("Unexpected end of event stream") + return None + buf += chunk + result = buf[:n] + buf = buf[n:] + return result + + while True: + if cancelled(): + return + # Prelude: total_length(4) + headers_length(4) + prelude_crc(4) + prelude_data = read_exactly(12) + if prelude_data is None: + return + + total_length, headers_length, _ = struct.unpack(">III", prelude_data) + + # total_length includes the 12-byte prelude. + remaining = total_length - 12 + if remaining <= 0: + continue + message_data = read_exactly(remaining) + if message_data is None: + return + + headers_bytes = message_data[:headers_length] + # Payload sits between headers and the trailing 4-byte message CRC. + payload_bytes = message_data[headers_length:-4] + + headers = {} + pos = 0 + while pos < len(headers_bytes): + name_len = headers_bytes[pos] + pos += 1 + if pos + name_len > len(headers_bytes): + break + name = headers_bytes[pos:pos + name_len].decode("utf-8") + pos += name_len + if pos >= len(headers_bytes): + break + header_type = headers_bytes[pos] + pos += 1 + if header_type == 7: # String type + if pos + 2 > len(headers_bytes): + break + value_len = struct.unpack( + ">H", headers_bytes[pos:pos + 2] + )[0] + pos += 2 + if pos + value_len > len(headers_bytes): + break + value = headers_bytes[pos:pos + value_len].decode("utf-8") + pos += value_len + headers[name] = value + else: + break + + message_type = headers.get(":message-type", "") + + if message_type == "exception": + error_msg = payload_bytes.decode("utf-8", errors="replace") + message, error_type = _parse_bedrock_error_body(error_msg) + raise BedrockHTTPError(0, message, error_type or "stream_error") + + if not payload_bytes: + continue + + event_type = headers.get(":event-type", "") + + try: + payload = json.loads(payload_bytes.decode("utf-8")) + except (json.JSONDecodeError, UnicodeDecodeError): + continue + + # Bedrock wraps the Anthropic event in {"bytes": ""} for + # chunk events. + if event_type == "chunk" and "bytes" in payload: + inner_bytes = base64.b64decode(payload["bytes"]) + try: + yield json.loads(inner_bytes.decode("utf-8")) + except (json.JSONDecodeError, UnicodeDecodeError): + continue + elif payload: + yield payload diff --git a/api/session_stats.py b/api/session_stats.py index d74cc4d..dfaf19e 100644 --- a/api/session_stats.py +++ b/api/session_stats.py @@ -38,6 +38,9 @@ def calculate_cost( price_tier = None model_lower = model.lower() + # Tier keys (haiku/sonnet/opus) are matched as case-insensitive + # substrings, so this also resolves Bedrock model ids like + # "us.anthropic.claude-sonnet-4-5-20250929-v1:0". for tier in pricing.keys(): if tier in model_lower: price_tier = pricing[tier] diff --git a/chat/ask_question.py b/chat/ask_question.py index a51c0f8..7d74ac8 100644 --- a/chat/ask_question.py +++ b/chat/ask_question.py @@ -77,31 +77,39 @@ def handle_input(self, code, question): if not self.create_chat_panel(): return - api_key = claudette_get_api_key_value() - - if not api_key: - window = self.get_window() - claudette_chat_status_message( - window, - ( - "Please add your Claude API key via the " - "`Settings > Package Settings > Claudette` menu." - ), - "⚠️", - ) - claudette_chat_status_message( - window, - ( - "Claudette allows you to define a single key, or you can " - "add multiple keys each with their own name. For example, " - 'you can define a "Work" and "Personal" key. If you have ' - "multiple API keys defined the " - "`Claudette: Switch API Key` command allows you switch " - "between them." - ), - "", - ) - return + self.load_settings() + provider = ( + self.settings.get("provider", "anthropic") + if self.settings + else "anthropic" + ) + + if provider != "bedrock": + api_key = claudette_get_api_key_value() + + if not api_key: + window = self.get_window() + claudette_chat_status_message( + window, + ( + "Please add your Claude API key via the " + "`Settings > Package Settings > Claudette` menu." + ), + "⚠️", + ) + claudette_chat_status_message( + window, + ( + "Claudette allows you to define a single key, or you " + "can add multiple keys each with their own name. For " + 'example, you can define a "Work" and "Personal" ' + "key. If you have multiple API keys defined the " + "`Claudette: Switch API Key` command allows you " + "switch between them." + ), + "", + ) + return self.send_to_claude(code, question.strip())