fix(llm-proxy): resolve $0 token costs and add Anthropic/OpenAI streaming (CAN-356)

canyonos_core/controller/cloud_provider_logic/EC2/_runtime.py

@@ -27,6 +27,9 @@
 from canyonos_core.controller.utils.env_file import env_file_args
 from canyonos_core.controller.utils.redis_utils import _wait_for_redis
 from canyonos_core.controller.utils.redis_client import RedisClient
+from canyonos_core.controller.cloud_provider_logic.shared_utils.llm_proxy_env import (
+    llm_proxy_docker_env_args,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -293,10 +296,8 @@ def _bootstrap_instance(
         f"CANYONOS_AGENT_PORT={CONTAINER_PORT}",
         "-e",
         f"CANYONOS_POLL_INTERVAL={_controller.config.get('poll_interval', 5)}",
-        # Route the agent's boto3 Bedrock calls through the in-container LLM
-        # proxy (started by LocalController) so token/cost telemetry is captured.
-        "-e",
-        "AWS_ENDPOINT_URL_BEDROCK_RUNTIME=http://127.0.0.1:8081/bedrock",
+        # Route the agent's LLM SDK calls through the in-container proxy for telemetry.
+        *llm_proxy_docker_env_args(),
         # The LLM stub is a local-only `canyonos test` control; it must never be
         # active on EC2. Pin it empty explicitly so a user's --env-file cannot
         # turn it on (docker: -e beats --env-file).

canyonos_core/controller/cloud_provider_logic/Local/_runtime.py

@@ -12,6 +12,9 @@
 
 from canyonos_core.controller.utils.container_names import container_name
 from canyonos_core.controller.utils.env_file import env_file_args
+from canyonos_core.controller.cloud_provider_logic.shared_utils.llm_proxy_env import (
+    llm_proxy_docker_env_args,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -116,10 +119,8 @@ def bootstrap_instance(provisioned, spec, replica_index, agent_id):
             f"CANYONOS_REDIS_PORT={spec.get('redis_port', 6379)}",
             "-e",
             f"CANYONOS_POLL_INTERVAL={_require_controller().config.get('poll_interval', 5)}",
-            # Route the agent's boto3 Bedrock calls through the in-container LLM
-            # proxy (started by LocalController) so token/cost telemetry is captured.
-            "-e",
-            "AWS_ENDPOINT_URL_BEDROCK_RUNTIME=http://127.0.0.1:8081/bedrock",
+            # Route the agent's LLM SDK calls through the in-container proxy for telemetry.
+            *llm_proxy_docker_env_args(),
         ]
 
         # LLM stub is a `canyonos test`-only control. `canyonos test` injects

canyonos_core/controller/cloud_provider_logic/shared_utils/llm_proxy_env.py

@@ -0,0 +1,47 @@
+"""Single source of truth for the env vars that route an agent's LLM SDK calls
+through the in-container LLM proxy (started by LocalController; see
+canyonos_core/llm_proxy/).
+
+Every LLM SDK/library reads a different env var name for its base URL -- there
+is no universal name to standardize on, so all of them still have to be set.
+What this module collapses is the VALUE: one PROXY_HOST constant, one place
+that lists the var names, used by both the Local and EC2 runtime backends
+instead of two independently hand-maintained copies of the same six lines.
+
+See company-memory: "LLM proxy can infer the provider from the request; the
+six base-URL variable names still cannot be collapsed" (CAN-343 analysis).
+"""
+
+# The proxy always runs on localhost inside the agent's own container.
+PROXY_HOST = "http://127.0.0.1:8081"
+
+
+def llm_proxy_env_vars(proxy_host: str = PROXY_HOST) -> dict:
+    """Return the {env_var_name: value} pairs that route Bedrock/OpenAI/Anthropic
+    SDK traffic through the in-container LLM proxy at `proxy_host`.
+    """
+    return {
+        # boto3 (Bedrock)
+        "AWS_ENDPOINT_URL_BEDROCK_RUNTIME": f"{proxy_host}/bedrock",
+        # openai SDK
+        "OPENAI_BASE_URL": f"{proxy_host}/openai/v1",
+        # langchain_openai / llama_index
+        "OPENAI_API_BASE": f"{proxy_host}/openai/v1",
+        # anthropic SDK
+        "ANTHROPIC_BASE_URL": f"{proxy_host}/anthropic",
+        # langchain_anthropic
+        "ANTHROPIC_API_URL": f"{proxy_host}/anthropic",
+        # LiteLLM
+        "ANTHROPIC_API_BASE": f"{proxy_host}/anthropic",
+    }
+
+
+def llm_proxy_docker_env_args(proxy_host: str = PROXY_HOST) -> list:
+    """Return `llm_proxy_env_vars` flattened into repeated `-e NAME=VALUE` pairs,
+    ready to splice into a `docker run` argument list.
+    """
+    args = []
+    for name, value in llm_proxy_env_vars(proxy_host).items():
+        args.append("-e")
+        args.append(f"{name}={value}")
+    return args

canyonos_core/controller/local_controller.py

@@ -84,6 +84,17 @@ def __init__(self, port=50051, publish_ready=True):
         redis_port = int(os.environ.get("CANYONOS_REDIS_PORT", 6379))
         self.redis = RedisClient(host=redis_host, port=redis_port)
         self._status_key = f"controller:{self.agent_host}:{self.public_port}:status"
+
+        # Every LLM call in this container is routed through the proxy, so it must be
+        # up before we report ready. The status key has no TTL: pin it to "failed" or
+        # a stale "healthy" keeps GlobalController seeing a container that has died.
+        try:
+            self._proxy_process = self._start_llm_proxy(redis_host, redis_port)
+        except Exception:
+            self.redis.set(self._status_key, "failed")
+            self.server.stop(0)
+            raise
+
         if publish_ready:
             self.redis.set(self._status_key, "healthy")
 
@@ -115,11 +126,6 @@ def __init__(self, port=50051, publish_ready=True):
         max_instances = int(os.environ.get("CANYONOS_MAX_AGENT_INSTANCES", 8))
         self._executor = ThreadPoolExecutor(max_workers=max_instances)
 
-        # Start the LLM proxy alongside the agent in this container. Bedrock
-        # calls are routed to it via AWS_ENDPOINT_URL_BEDROCK_RUNTIME (injected
-        # by the runtime), and it writes token/cost telemetry to Redis.
-        self._proxy_process = self._start_llm_proxy(redis_host, redis_port)
-
         logger.info(
             "Local controller initialized at %s (max_agent_instances=%d), reported healthy to Redis.",
             self._my_endpoint,
@@ -138,31 +144,78 @@ def mark_failed(self):
     def _start_llm_proxy(self, redis_host, redis_port):
         """Start the LLM proxy as a subprocess in this container (127.0.0.1:8081).
 
-        Best-effort: a failure here must never stop the controller from coming up.
+        Fatal: the runtime force-injects LLM base-URL env vars that point every
+        Bedrock/OpenAI/Anthropic SDK call at this proxy (`docker run -e` beats
+        `--env-file`, so nothing can opt back out). If it fails to start, every
+        LLM call in this container would silently fail or hang, not just lose
+        telemetry -- so raise instead of limping on with no proxy listening.
         """
+        import socket
         import subprocess
 
+        import requests
+
+        # An orphaned proxy on 8081 would answer the /healthz probe below and mask
+        # one of ours that never bound, so prove the port free before spawning.
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as port_probe:
+            port_probe.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+            try:
+                port_probe.bind(("127.0.0.1", 8081))
+            except OSError as e:
+                logger.error("127.0.0.1:8081 is already in use: %s", e)
+                raise RuntimeError(
+                    "127.0.0.1:8081 is already in use, so the LLM proxy cannot bind "
+                    "it; agent LLM calls are routed through it unconditionally and "
+                    "would otherwise fail silently."
+                ) from e
+
+        proxy_env = os.environ.copy()
+        proxy_env.update(
+            {
+                "PROXY_HOST": "127.0.0.1",
+                "PROXY_PORT": "8081",
+                "CANYONOS_REDIS_HOST": redis_host,
+                "CANYONOS_REDIS_PORT": str(redis_port),
+            }
+        )
         try:
-            proxy_env = os.environ.copy()
-            proxy_env.update(
-                {
-                    "PROXY_HOST": "127.0.0.1",
-                    "PROXY_PORT": "8081",
-                    "CANYONOS_REDIS_HOST": redis_host,
-                    "CANYONOS_REDIS_PORT": str(redis_port),
-                }
-            )
             proxy_process = subprocess.Popen(
                 [sys.executable, "-m", "canyonos_core.llm_proxy"],
                 env=proxy_env,
             )
-            logger.info(
-                "Started LLM proxy on 127.0.0.1:8081 (PID: %d)", proxy_process.pid
-            )
-            return proxy_process
         except Exception as e:
-            logger.warning("Failed to start LLM proxy: %s", e)
-            return None
+            logger.error("Failed to start LLM proxy: %s", e)
+            raise RuntimeError(
+                "LLM proxy failed to start; agent LLM calls are routed through it "
+                "unconditionally and would otherwise fail silently."
+            ) from e
+
+        # Popen only raises if the process can't be spawned -- it returns a healthy
+        # handle even if the proxy starts and dies immediately, so poll /healthz.
+        deadline = time.time() + 10
+        while time.time() < deadline:
+            if proxy_process.poll() is not None:
+                raise RuntimeError(
+                    f"LLM proxy exited immediately (code {proxy_process.returncode}); "
+                    "agent LLM calls are routed through it unconditionally and would "
+                    "otherwise fail silently."
+                )
+            try:
+                if requests.get("http://127.0.0.1:8081/healthz", timeout=0.5).ok:
+                    break
+            except requests.exceptions.RequestException:
+                pass
+            time.sleep(0.2)
+        else:
+            proxy_process.kill()
+            raise RuntimeError(
+                "LLM proxy did not become healthy on 127.0.0.1:8081 within 10s; "
+                "agent LLM calls are routed through it unconditionally and would "
+                "otherwise fail silently."
+            )
+
+        logger.info("Started LLM proxy on 127.0.0.1:8081 (PID: %d)", proxy_process.pid)
+        return proxy_process
 
     def _collect_metrics(self):
         """Snapshot current instance health/resource metrics.

canyonos_core/controller/utils/pricing.py

@@ -1,10 +1,11 @@
-"""Lazily-cached lookups over the static aws_pricing_chart.db reference data."""
+"""Lazily-cached lookups over the static llm_token_costs.db reference data."""
 
 import os
+import re
 
 from sqlalchemy import create_engine, text
 
-_PRICING_DB_PATH = os.path.join(os.path.dirname(__file__), "aws_pricing_chart.db")
+_PRICING_DB_PATH = os.path.join(os.path.dirname(__file__), "llm_token_costs.db")
 
 _hourly_cost_by_instance_type: dict[str, float] | None = None
 _token_cost_by_model_id: dict[str, tuple[float, float]] | None = None
@@ -23,7 +24,7 @@ def _load_cache():
         model_rows = conn.execute(
             text(
                 "SELECT model_id, input_cost_per_million_tokens, "
-                "output_cost_per_million_tokens FROM bedrock_model_pricing"
+                "output_cost_per_million_tokens FROM llm_token_costs"
             )
         ).fetchall()
     engine.dispose()
@@ -33,18 +34,28 @@ def _load_cache():
 
 
 def compute_token_cost(model_id, input_token_count, output_token_count):
-    """Return the USD cost of a Bedrock call, or 0.0 if the model_id is unknown."""
+    """Return the USD cost of an LLM call, or 0.0 if the model_id is unknown."""
     _load_cache()
-    if _token_cost_by_model_id is None:
-        return 0.0
-    costs = _token_cost_by_model_id.get(model_id)
-    if costs is None:
-        return 0.0
-    input_cost_per_million, output_cost_per_million = costs
-    return (
-        input_token_count * input_cost_per_million
-        + output_token_count * output_cost_per_million
-    ) / 1_000_000
+    if _token_cost_by_model_id and isinstance(model_id, str) and model_id:
+        candidates = [model_id]
+        undated = re.sub(r"-\d{4}-?\d{2}-?\d{2}$", "", model_id)
+        if undated != model_id:
+            candidates.append(undated)
+        if model_id.startswith("claude-"):
+            # Direct-API Anthropic ids carry no vendor prefix or version suffix, and the
+            # table is inconsistent about the date segment, so try both spellings.
+            candidates.append(f"anthropic.{model_id}-v1:0")
+            if undated != model_id:
+                candidates.append(f"anthropic.{undated}-v1:0")
+        for candidate in candidates:
+            costs = _token_cost_by_model_id.get(candidate)
+            if costs is not None:
+                input_cost_per_million, output_cost_per_million = costs
+                return (
+                    input_token_count * input_cost_per_million
+                    + output_token_count * output_cost_per_million
+                ) / 1_000_000
+    return 0.0
 
 
 def compute_server_cost(instance_type, execution_time_seconds):

canyonos_core/llm_proxy/README.md

@@ -5,8 +5,7 @@ A local, single-machine pass-through proxy for **OpenAI**, **Anthropic**, and
 one base-URL env var per provider. Every call flows through one function
 (`core.proxy_request`) where token/metrics hooks fire.
 
-**Scope:** request/response ("call and return") only. Streaming is intentionally
-not implemented yet.
+**Scope:** request/response and streaming, for all three providers.
 
 ## How it works
 
@@ -19,7 +18,10 @@ your app (unchanged)         localhost:8080                 real upstream
 ```
 
 - **OpenAI / Anthropic** — straight HTTP reverse-proxy: rewrite host, swap in the
-  real key, forward with `requests`, return the response.
+  real key, forward with `requests`, return the response. Server-sent-event
+  responses are relayed byte-for-byte while usage is folded out of the events. An
+  upstream failure mid-stream aborts the response rather than ending it cleanly,
+  so a partial answer can't reach the caller looking like a complete one.
 - **Bedrock** — re-issued through the proxy's own `boto3` client (handles SigV4
   signing + URL-encoding correctly). `invoke`, `converse`, `converse-stream`,
   and `invoke-with-response-stream` are all wired up.
@@ -95,33 +97,30 @@ Telemetry is written to Redis under `future:<future_id>` keys, keyed off an
 | Bedrock `converse` / `converse-stream` | Bedrock's own camelCase (`inputTokens`, ...) | works for any model |
 | Bedrock `invoke` / `invoke-with-response-stream`, `anthropic.*` model | Anthropic's native (`input_tokens`, ...) | works |
 | Bedrock `invoke` / `invoke-with-response-stream`, other model families | model-specific, unknown | no usage (schema not implemented yet) |
-| Direct Anthropic API (`/anthropic/...`) | Anthropic's native | works |
-| Direct OpenAI API (`/openai/...`) | OpenAI's native (`prompt_tokens`, ...) | works |
+| Direct Anthropic API (`/anthropic/...`) | Anthropic's native | works, streaming and non-streaming |
+| Direct OpenAI API (`/openai/...`) | OpenAI's native (`prompt_tokens`, ...) | works; streaming needs the caller to set `stream_options.include_usage` |
 
 ### How it works
 
-1. **Auto-injection:** boto3 hook (`proxy.py`) injects `X-Canyonos-Future-ID` header from thread-local context -- **Bedrock (boto3) only**; the OpenAI/Anthropic SDKs don't fire this hook, so callers using those SDKs directly won't get the header auto-attached.
+1. **Auto-injection:** `proxy.py` injects the `X-Canyonos-Future-ID` header from thread-local context for all three providers -- a boto3 event hook for Bedrock, and an `httpx.Client.send` patch for the OpenAI/Anthropic SDKs, gated to proxy-bound paths.
 2. **Usage extraction:** `hooks.py`'s `_extract_usage` parses the response `usage` field with the schema matching that provider/op/model (table above) -- this part works for all three providers whenever the header is present.
 3. **Redis write:** All metrics written to `future:<future_id>` hash.
 
-So in practice, automatic end-to-end telemetry (header injection + extraction) is Bedrock-only for now; OpenAI/Anthropic usage extraction works, but nothing auto-attaches `X-Canyonos-Future-ID` for those SDKs yet.
-
-### Why is header auto-injection Bedrock-only?
-
-OpenAI and Anthropic use their own Python SDKs (`openai`, `anthropic`), not boto3.
-The boto3 event hook doesn't fire for non-AWS SDKs. To add auto-injection for those:
-- Would need separate hooks in each SDK's HTTP client
-- Or callers would need to attach `X-Canyonos-Future-ID` themselves
-
 ## Limitations
 
-- **Bedrock `converse-stream` and `invoke-with-response-stream` only.**
-  OpenAI/Anthropic `stream=True` is still not handled and remains fully
-  buffered. `invoke-with-response-stream` also has no usage/token telemetry
-  regardless of model (unlike `converse-stream`, it has no metadata event to
-  read usage from -- see the usage extraction coverage table above).
+- **OpenAI streaming reports usage only when the caller sets
+  `stream_options: {"include_usage": true}`** -- OpenAI omits usage from the
+  stream otherwise, and the proxy does not rewrite the caller's request body.
+- **Bedrock `invoke-with-response-stream` has no usage/token telemetry**
+  regardless of model: unlike `converse-stream`, it has no metadata event to
+  read usage from (see the usage extraction coverage table above).
 - **Bedrock error bodies are reconstructed**, not passed through byte-for-byte
   (boto3 raises on 4xx/5xx; we rebuild a JSON body with the real status +
   message). OpenAI/Anthropic errors pass through unchanged.
 - **Dev server.** Runs on Flask's built-in server — fine for a local proxy, not
   meant for production traffic.
+- **Agents can't point at a custom OpenAI/Anthropic-compatible endpoint
+  themselves** (Azure OpenAI, a self-hosted vLLM/Ollama, OpenRouter, ...) --
+  `OPENAI_BASE_URL`/`ANTHROPIC_BASE_URL`/etc. are force-pinned at the proxy, so
+  a value the agent sets is ignored. Use `OPENAI_UPSTREAM_BASE` /
+  `ANTHROPIC_UPSTREAM_BASE` on the proxy process instead to route it there.

canyonos_core/llm_proxy/hooks.py

@@ -3,8 +3,7 @@
 Every proxied call passes through ``on_request`` / ``on_response``, which logs
 and (when Redis is configured) extracts token usage per provider/op/model --
 see ``Hooks._extract_usage``. Bedrock invoke's usage schema is only known for
-anthropic.* models today; other model families and OpenAI/Anthropic streaming
-remain unhandled.
+anthropic.* models today; other model families remain unhandled.
 """
 
 from __future__ import annotations
@@ -178,9 +177,15 @@ def _extract_usage(self, ctx: Ctx, resp: Any) -> Optional[TokenUsage]:
             return self._extract_json_usage(resp, self._usage_from_dict)
 
         if ctx.provider == "anthropic":
+            if is_stream:
+                return self._usage_from_anthropic_dict(
+                    getattr(resp, "stream_usage", None)
+                )
             return self._extract_json_usage(resp, self._usage_from_anthropic_dict)
 
         if ctx.provider == "openai":
+            if is_stream:
+                return self._usage_from_openai_dict(getattr(resp, "stream_usage", None))
             return self._extract_json_usage(resp, self._usage_from_openai_dict)
 
         return None

canyonos_core/llm_proxy/providers/anthropic.py

@@ -21,3 +21,15 @@ def target(self, req, subpath, body):
             headers=headers,
             params=req.args.to_dict(flat=True),
         )
+
+    def merge_stream_usage(self, payload, usage):
+        # Input counts arrive on message_start and output counts on message_delta, so neither event alone is enough.
+        if payload.get("type") == "message_start":
+            incoming = (payload.get("message") or {}).get("usage") or {}
+        elif payload.get("type") == "message_delta":
+            incoming = payload.get("usage") or {}
+        else:
+            return
+        for key, value in incoming.items():
+            if value:
+                usage[key] = value