From 29a4ca028dcaf3ecf5cd5971474414cf6d01b915 Mon Sep 17 00:00:00 2001 From: Matt Bourke Date: Tue, 12 May 2026 14:10:45 +1000 Subject: [PATCH 1/2] Add AWS Bedrock provider support - Allow routing requests through AWS Bedrock with SigV4 authentication as an alternative to the direct Anthropic API. --- Claudette.sublime-settings | 16 ++ api/api.py | 218 +++++++++++++------------ api/bedrock.py | 319 +++++++++++++++++++++++++++++++++++++ chat/ask_question.py | 58 ++++--- 4 files changed, 487 insertions(+), 124 deletions(-) create mode 100644 api/bedrock.py diff --git a/Claudette.sublime-settings b/Claudette.sublime-settings index 49e4e31..e811d5a 100644 --- a/Claudette.sublime-settings +++ b/Claudette.sublime-settings @@ -1,4 +1,20 @@ { + // 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` + // "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..1d1b208 100644 --- a/api/api.py +++ b/api/api.py @@ -25,6 +25,12 @@ ) from ..utils import claudette_chat_status_message, claudette_get_api_key_value from . import session_stats +from .bedrock import ( + BEDROCK_ANTHROPIC_VERSION, + bedrock_request, + get_aws_credentials, + parse_event_stream, +) from .cancellation import CancellationToken from .errors import ( handle_model_not_found, @@ -49,8 +55,10 @@ class CancelledException(Exception): class ClaudetteClaudeAPI: def __init__(self): self.settings = sublime.load_settings(SETTINGS_FILE) + self.provider = self.settings.get("provider", "anthropic") 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 +72,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 +237,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 +345,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." @@ -638,7 +679,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 +689,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 +698,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 +706,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 +718,56 @@ 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 - try: - # Try to get the underlying socket - if hasattr(response.fp, "raw"): - raw = response.fp.raw - if hasattr(raw, "_sock"): - sock = raw._sock - 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, - ) + def _iter_events_sse(response): + """Yield parsed JSON dicts from Anthropic SSE stream.""" + for line in response: + if not line or line.isspace(): + continue + chunk = line.decode("utf-8") + if not chunk.startswith("data: "): + continue + chunk = chunk[6:] + if chunk.strip() == "[DONE]": 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 - - if line.isspace(): + yield json.loads(chunk) + except (json.JSONDecodeError, ValueError): continue - try: - chunk = line.decode("utf-8") - if not chunk.startswith("data: "): - continue + 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) + 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) - chunk = chunk[6:] # Remove 'data: ' prefix - if chunk.strip() == "[DONE]": - break + response, event_iter = _open_and_iter() + try: + for data in event_iter: + if is_cancelled(): + break - data = json.loads(chunk) + try: # Get initial input tokens from message_start if data.get("type") == "message_start": @@ -1048,8 +1053,10 @@ def _send_citation( ) except Exception: - # Skip invalid chunks without error messages continue + finally: + if hasattr(response, "close"): + response.close() except urllib.error.HTTPError as e: error_type, error_message = parse_api_error(e) @@ -1061,16 +1068,29 @@ 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 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: + return [ + "anthropic.claude-sonnet-4-6-20250514-v1:0", + "anthropic.claude-opus-4-7-20250514-v1:0", + "anthropic.claude-sonnet-4-5-20241022-v2:0", + "anthropic.claude-haiku-4-5-20241022-v1:0", + "us.anthropic.claude-sonnet-4-6-20250514-v1:0", + "us.anthropic.claude-sonnet-4-5-20241022-v2:0", + "us.anthropic.claude-haiku-4-5-20241022-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..c82c2fb --- /dev/null +++ b/api/bedrock.py @@ -0,0 +1,319 @@ +import base64 +import hashlib +import hmac +import http.client +import json +import os +import struct +import subprocess +import ssl +import urllib.parse +from datetime import datetime, timezone + +BEDROCK_ANTHROPIC_VERSION = 'bedrock-2023-05-31' + + +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): + """Double-encode the path for SigV4 canonical request.""" + segments = path.split('/') + return '/'.join(urllib.parse.quote(seg, safe='') for seg in segments) + + +def _get_credentials_from_profile(profile=None): + """Get AWS credentials from CLI profile using 'aws configure export-credentials'.""" + cmd = ['aws', 'configure', 'export-credentials', '--format', 'env'] + if profile: + cmd.extend(['--profile', profile]) + try: + result = subprocess.run( + cmd, + capture_output=True, + text=True, + timeout=10 + ) + if result.returncode != 0: + return None + creds = {} + for line in result.stdout.strip().split('\n'): + if '=' in line: + line = line.replace('export ', '') + key, _, value = line.partition('=') + creds[key.strip()] = 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', '') + if access_key and secret_key: + return { + 'access_key': access_key, + 'secret_key': secret_key, + 'session_token': session_token + } + except (subprocess.TimeoutExpired, FileNotFoundError, OSError): + pass + return None + + +def _get_credentials_from_env(): + """Get AWS credentials from environment variables.""" + 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', '') + if access_key and secret_key: + return { + 'access_key': access_key, + 'secret_key': secret_key, + 'session_token': session_token + } + return None + + +def get_aws_credentials(settings): + """Resolve AWS credentials from settings, env vars, or AWS CLI profile.""" + access_key = settings.get('aws_access_key_id', '') + secret_key = settings.get('aws_secret_access_key', '') + session_token = settings.get('aws_session_token', '') + if access_key and secret_key: + return { + 'access_key': access_key, + 'secret_key': secret_key, + 'session_token': session_token + } + creds = _get_credentials_from_env() + if creds: + return creds + profile = settings.get('aws_profile', '') + return _get_credentials_from_profile(profile or 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 '/model/{0}/{1}'.format(encoded_model, endpoint) + + +def bedrock_request(region, model_id, body_dict, credentials, streaming=False, verify_ssl=True): + """ + Make a signed request to AWS Bedrock and return the response. + For non-streaming: returns parsed JSON dict. + For streaming: returns the http.client.HTTPResponse (caller must read and close). + """ + host = 'bedrock-runtime.{0}.amazonaws.com'.format(region) + 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() + + # Headers to sign + 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( + '{0}:{1}\n'.format(k, signed_headers_dict[k]) for k in signed_header_keys + ) + + # SigV4 requires double-encoding the path in the canonical request + 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 = '{0}/{1}/{2}/aws4_request'.format(date_stamp, region, service) + 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 = ( + 'AWS4-HMAC-SHA256 Credential={0}/{1}, SignedHeaders={2}, Signature={3}' + ).format(credentials['access_key'], credential_scope, signed_headers_str, signature) + + # Build final headers for the HTTP request + 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'] + + # Use http.client to avoid urllib's URL re-encoding + if verify_ssl: + context = ssl.create_default_context() + else: + context = ssl.create_default_context() + context.check_hostname = False + context.verify_mode = ssl.CERT_NONE + + conn = http.client.HTTPSConnection(host, context=context) + conn.request('POST', request_path, body=body, headers=headers) + response = conn.getresponse() + + if response.status != 200: + error_body = response.read().decode('utf-8', errors='replace') + conn.close() + try: + error_data = json.loads(error_body) + msg = error_data.get('message', error_body) + except (json.JSONDecodeError, KeyError): + msg = error_body + raise RuntimeError("Bedrock HTTP {0}: {1}".format(response.status, msg)) + + if not streaming: + raw = response.read().decode('utf-8') + conn.close() + return json.loads(raw) + + # For streaming, return the response object — caller manages reading + # Attach the connection so caller can close it + response._conn = conn + return response + + +def parse_event_stream(response): + """ + Parse AWS event stream binary format from a Bedrock streaming response. + Yields parsed JSON dicts matching the Anthropic SSE event format. + """ + buf = b'' + + def read_exactly(n): + nonlocal buf + while len(buf) < n: + chunk = response.read(n - len(buf)) + 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: + # Read 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) + + # Read rest of message (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 + + # Parse headers + headers_bytes = message_data[:headers_length] + # Payload is between headers and message CRC (last 4 bytes) + payload_bytes = message_data[headers_length:-4] + + # Parse event stream headers (name-value pairs with type byte) + headers = {} + pos = 0 + while pos < len(headers_bytes): + if pos >= len(headers_bytes): + break + 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 + + # Check for exceptions + message_type = headers.get(':message-type', '') + + if message_type == 'exception': + error_msg = payload_bytes.decode('utf-8', errors='replace') + try: + error_data = json.loads(error_msg) + raise RuntimeError("Bedrock stream error: {0}".format( + error_data.get('message', error_msg) + )) + except json.JSONDecodeError: + raise RuntimeError("Bedrock stream error: {0}".format(error_msg)) + + 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/chat/ask_question.py b/chat/ask_question.py index a51c0f8..3d68dc6 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()) From beac4f882cf865fbe7f956d89519c0fd5fe94946 Mon Sep 17 00:00:00 2001 From: Matt Bourke Date: Fri, 12 Jun 2026 16:52:40 +1000 Subject: [PATCH 2/2] Address PR review feedback for AWS Bedrock provider Fix streaming cancellation regression on the Anthropic path: restore non-blocking SSE polling so cancellation can fire every ~500ms, and ensure cancellation triggers chunk_callback("", is_done=True, was_cancelled=True) instead of leaving the response heading dangling. Add equivalent cooperative cancellation to the Bedrock event-stream parser via an optional should_cancel callback. Bedrock hardening: pass timeout=30 to HTTPSConnection (was unbounded); introduce BedrockHTTPError(status, message, error_type) so 4xx responses route through the existing model-not-found UX; cache resolved AWS credentials per source (settings/env/profile) honoring their reported Expiration; suppress the AWS CLI console window on Windows; tighten 'aws configure export-credentials' parsing to use startswith('export ') and strip surrounding quotes; drop the fragile response._conn attribute in favor of a small wrapper that owns both the response and connection. Replace fabricated Bedrock model ids with verified shipping ones (Opus 4.6/4.5, Sonnet 4.6/4.5, Haiku 4.5) plus their us./eu./apac./ global. inference-profile equivalents. Validate the provider setting and warn-and-fall-back on unknown values rather than silently routing Bedrock-style ids to the Anthropic endpoint. Normalize api/bedrock.py to project conventions: double quotes, f-strings, line length, alphabetical imports, Google-style docstrings on public functions. Document the AWS CLI requirement and credential caching in the settings file. Revert the unrelated re-quoting in chat/ask_question.py. Co-Authored-By: Claude Opus 4.7 (1M context) --- Claudette.sublime-settings | 3 + api/api.py | 108 ++++++- api/bedrock.py | 579 +++++++++++++++++++++++++++---------- api/session_stats.py | 3 + chat/ask_question.py | 2 +- 5 files changed, 527 insertions(+), 168 deletions(-) diff --git a/Claudette.sublime-settings b/Claudette.sublime-settings index e811d5a..66abd70 100644 --- a/Claudette.sublime-settings +++ b/Claudette.sublime-settings @@ -10,6 +10,9 @@ // 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": "", diff --git a/api/api.py b/api/api.py index 1d1b208..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 @@ -27,11 +26,11 @@ from . import session_stats from .bedrock import ( BEDROCK_ANTHROPIC_VERSION, + BedrockHTTPError, bedrock_request, get_aws_credentials, parse_event_stream, ) -from .cancellation import CancellationToken from .errors import ( handle_model_not_found, is_model_not_found_error, @@ -45,6 +44,8 @@ parse_web_search_items, ) +KNOWN_PROVIDERS = ("anthropic", "bedrock") + class CancelledException(Exception): """Raised when a request is cancelled.""" @@ -55,7 +56,16 @@ class CancelledException(Exception): class ClaudetteClaudeAPI: def __init__(self): self.settings = sublime.load_settings(SETTINGS_FILE) - self.provider = self.settings.get("provider", "anthropic") + 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") @@ -406,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: @@ -719,9 +744,26 @@ def is_cancelled(): stream_current_block_index = None def _iter_events_sse(response): - """Yield parsed JSON dicts from Anthropic SSE stream.""" - for line in response: - if not line or line.isspace(): + """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: + response.fp._sock.settimeout(0.5) + except Exception: + pass + while True: + if is_cancelled(): + return + try: + line = response.readline() + except socket.timeout: + continue + if not line: + return + if line.isspace(): continue chunk = line.decode("utf-8") if not chunk.startswith("data: "): @@ -742,7 +784,9 @@ def _open_and_iter(): self.aws_region, self.model, data, credentials, streaming=True, verify_ssl=self.verify_ssl, ) - return resp, parse_event_stream(resp) + return resp, parse_event_stream( + resp, should_cancel=is_cancelled + ) else: headers = { "x-api-key": self.api_key, @@ -1054,6 +1098,17 @@ def _send_citation( except Exception: 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() @@ -1069,6 +1124,16 @@ def _send_citation( handle_error("[Error] {0}".format(error_message)) except urllib.error.URLError as 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: @@ -1081,14 +1146,29 @@ def _send_citation( 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-sonnet-4-6-20250514-v1:0", - "anthropic.claude-opus-4-7-20250514-v1:0", - "anthropic.claude-sonnet-4-5-20241022-v2:0", - "anthropic.claude-haiku-4-5-20241022-v1:0", - "us.anthropic.claude-sonnet-4-6-20250514-v1:0", - "us.anthropic.claude-sonnet-4-5-20241022-v2:0", - "us.anthropic.claude-haiku-4-5-20241022-v1:0", + "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: diff --git a/api/bedrock.py b/api/bedrock.py index c82c2fb..6b0c762 100644 --- a/api/bedrock.py +++ b/api/bedrock.py @@ -1,188 +1,429 @@ +"""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 ssl +import sys +import threading import urllib.parse -from datetime import datetime, timezone +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) -BEDROCK_ANTHROPIC_VERSION = 'bedrock-2023-05-31' + +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') + msg = msg.encode("utf-8") if isinstance(key, str): - key = key.encode('utf-8') + 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_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') + k_signing = _hmac_sha256(k_service, "aws4_request") return k_signing def _uri_encode_path_for_signing(path): - """Double-encode the path for SigV4 canonical request.""" - segments = path.split('/') - return '/'.join(urllib.parse.quote(seg, safe='') for seg in segments) + """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): - """Get AWS credentials from CLI profile using 'aws configure export-credentials'.""" - cmd = ['aws', 'configure', 'export-credentials', '--format', 'env'] + """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]) + cmd.extend(["--profile", profile]) try: - result = subprocess.run( - cmd, - capture_output=True, - text=True, - timeout=10 - ) - if result.returncode != 0: - return None - creds = {} - for line in result.stdout.strip().split('\n'): - if '=' in line: - line = line.replace('export ', '') - key, _, value = line.partition('=') - creds[key.strip()] = 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', '') - if access_key and secret_key: - return { - 'access_key': access_key, - 'secret_key': secret_key, - 'session_token': session_token - } + result = subprocess.run(cmd, **_subprocess_kwargs_no_window()) except (subprocess.TimeoutExpired, FileNotFoundError, OSError): - pass + 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_env(): - """Get AWS credentials from environment variables.""" - 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', '') +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 + "access_key": access_key, + "secret_key": secret_key, + "session_token": session_token, + "expiration": None, } return None -def get_aws_credentials(settings): - """Resolve AWS credentials from settings, env vars, or AWS CLI profile.""" - access_key = settings.get('aws_access_key_id', '') - secret_key = settings.get('aws_secret_access_key', '') - session_token = settings.get('aws_session_token', '') +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 + "access_key": access_key, + "secret_key": secret_key, + "session_token": session_token, + "expiration": None, } - creds = _get_credentials_from_env() - if creds: - return creds - profile = settings.get('aws_profile', '') - return _get_credentials_from_profile(profile or 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 '/model/{0}/{1}'.format(encoded_model, 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 bedrock_request(region, model_id, body_dict, credentials, streaming=False, verify_ssl=True): +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. """ - Make a signed request to AWS Bedrock and return the response. - For non-streaming: returns parsed JSON dict. - For streaming: returns the http.client.HTTPResponse (caller must read and close). + + 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 = 'bedrock-runtime.{0}.amazonaws.com'.format(region) + host = f"bedrock-runtime.{region}.amazonaws.com" request_path = _build_request_path(model_id, streaming=streaming) - body = json.dumps(body_dict).encode('utf-8') + 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') + amz_date = now.strftime("%Y%m%dT%H%M%SZ") + date_stamp = now.strftime("%Y%m%d") payload_hash = hashlib.sha256(body).hexdigest() - # Headers to sign signed_headers_dict = { - 'content-type': 'application/json', - 'host': host, - 'x-amz-content-sha256': payload_hash, - 'x-amz-date': amz_date, + "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'] + 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( - '{0}:{1}\n'.format(k, signed_headers_dict[k]) for k in signed_header_keys + signed_headers_str = ";".join(signed_header_keys) + canonical_headers = "".join( + f"{k}:{signed_headers_dict[k]}\n" for k in signed_header_keys ) - # SigV4 requires double-encoding the path in the canonical request canonical_uri = _uri_encode_path_for_signing(request_path) - service = 'bedrock' + service = "bedrock" - canonical_request = '\n'.join([ - 'POST', + canonical_request = "\n".join([ + "POST", canonical_uri, - '', # empty query string + "", # empty query string canonical_headers, signed_headers_str, - payload_hash + payload_hash, ]) - credential_scope = '{0}/{1}/{2}/aws4_request'.format(date_stamp, region, service) - string_to_sign = '\n'.join([ - 'AWS4-HMAC-SHA256', + 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() + hashlib.sha256(canonical_request.encode("utf-8")).hexdigest(), ]) signing_key = _get_signature_key( - credentials['secret_key'], date_stamp, region, service + credentials["secret_key"], date_stamp, region, service ) signature = hmac.new( - signing_key, string_to_sign.encode('utf-8'), hashlib.sha256 + signing_key, string_to_sign.encode("utf-8"), hashlib.sha256 ).hexdigest() authorization = ( - 'AWS4-HMAC-SHA256 Credential={0}/{1}, SignedHeaders={2}, Signature={3}' - ).format(credentials['access_key'], credential_scope, signed_headers_str, signature) + f"AWS4-HMAC-SHA256 Credential={credentials['access_key']}/" + f"{credential_scope}, SignedHeaders={signed_headers_str}, " + f"Signature={signature}" + ) - # Build final headers for the HTTP request headers = { - 'Content-Type': 'application/json', - 'Host': host, - 'X-Amz-Date': amz_date, - 'X-Amz-Content-Sha256': payload_hash, - 'Authorization': authorization, + "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 credentials.get("session_token"): + headers["X-Amz-Security-Token"] = credentials["session_token"] - # Use http.client to avoid urllib's URL re-encoding if verify_ssl: context = ssl.create_default_context() else: @@ -190,42 +431,79 @@ def bedrock_request(region, model_id, body_dict, credentials, streaming=False, v context.check_hostname = False context.verify_mode = ssl.CERT_NONE - conn = http.client.HTTPSConnection(host, context=context) - conn.request('POST', request_path, body=body, headers=headers) - response = conn.getresponse() + # 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') + error_body = response.read().decode("utf-8", errors="replace") conn.close() - try: - error_data = json.loads(error_body) - msg = error_data.get('message', error_body) - except (json.JSONDecodeError, KeyError): - msg = error_body - raise RuntimeError("Bedrock HTTP {0}: {1}".format(response.status, msg)) + message, error_type = _parse_bedrock_error_body(error_body) + raise BedrockHTTPError(response.status, message, error_type) if not streaming: - raw = response.read().decode('utf-8') - conn.close() - return json.loads(raw) + try: + raw = response.read().decode("utf-8") + return json.loads(raw) + finally: + conn.close() - # For streaming, return the response object — caller manages reading - # Attach the connection so caller can close it - response._conn = conn - return response + return _BedrockStreamResponse(response, conn) -def parse_event_stream(response): - """ - Parse AWS event stream binary format from a Bedrock streaming response. - Yields parsed JSON dicts matching the Anthropic SSE event format. +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'' + 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: - chunk = response.read(n - len(buf)) + 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") @@ -236,14 +514,16 @@ def read_exactly(n): return result while True: - # Read prelude: total_length(4) + headers_length(4) + prelude_crc(4) + 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, headers_length, _ = struct.unpack(">III", prelude_data) - # Read rest of message (total_length includes the 12-byte prelude) + # total_length includes the 12-byte prelude. remaining = total_length - 12 if remaining <= 0: continue @@ -251,22 +531,18 @@ def read_exactly(n): if message_data is None: return - # Parse headers headers_bytes = message_data[:headers_length] - # Payload is between headers and message CRC (last 4 bytes) + # Payload sits between headers and the trailing 4-byte message CRC. payload_bytes = message_data[headers_length:-4] - # Parse event stream headers (name-value pairs with type byte) headers = {} pos = 0 while pos < len(headers_bytes): - if pos >= len(headers_bytes): - break 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') + name = headers_bytes[pos:pos + name_len].decode("utf-8") pos += name_len if pos >= len(headers_bytes): break @@ -275,44 +551,41 @@ def read_exactly(n): if header_type == 7: # String type if pos + 2 > len(headers_bytes): break - value_len = struct.unpack('>H', headers_bytes[pos:pos + 2])[0] + 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') + value = headers_bytes[pos:pos + value_len].decode("utf-8") pos += value_len headers[name] = value else: break - # Check for exceptions - message_type = headers.get(':message-type', '') + message_type = headers.get(":message-type", "") - if message_type == 'exception': - error_msg = payload_bytes.decode('utf-8', errors='replace') - try: - error_data = json.loads(error_msg) - raise RuntimeError("Bedrock stream error: {0}".format( - error_data.get('message', error_msg) - )) - except json.JSONDecodeError: - raise RuntimeError("Bedrock stream error: {0}".format(error_msg)) + 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', '') + event_type = headers.get(":event-type", "") try: - payload = json.loads(payload_bytes.decode('utf-8')) + 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']) + # 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')) + yield json.loads(inner_bytes.decode("utf-8")) except (json.JSONDecodeError, UnicodeDecodeError): continue elif 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 3d68dc6..7d74ac8 100644 --- a/chat/ask_question.py +++ b/chat/ask_question.py @@ -102,7 +102,7 @@ def handle_input(self, code, question): ( "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\" " + '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."