v0.14.1 apiproxy: auto-resolve omc_proxy_expand_ref tool calls transparently

When the upstream response contains a sole tool_use for omc_proxy_expand_ref,
the proxy now:
  1. Pulls the hash_str from the tool_use input
  2. Looks it up in the MemoryStore cache
  3. Builds a follow-up request: original messages + assistant response +
     synthetic user turn with tool_result containing the expanded text
  4. Re-calls upstream (bounded to 8 rounds to prevent runaway)
  5. Returns the final non-expand-call response to the client

The client never sees the omc_proxy_expand_ref round-trip. Solves the
v0.14.0-alpha sharp edge where Claude Code would otherwise see the
unknown expand tool and error.

End-to-end smoke test with a mock upstream that simulates the LLM asking
to expand: client sends 7KB request, proxy rewrites to 908 bytes upstream,
mock emits expand tool_use, proxy resolves from cache, mock sees 6.8KB
tool_result, mock emits final text, client receives a single clean
assistant message. Round-trip lossless.

Mixed tool_use (expand alongside another tool) still passes through —
v0.14.3 follow-up will handle that case. Streaming still v0.14.2.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Author: RandomCoder-lab

omnimcode-apiproxy/README.md

@@ -2,7 +2,7 @@
 
 Substrate-rewriting reverse proxy for `api.anthropic.com`. Compresses large content blocks in the LLM's context window by replacing them with content-addressed `<omc:ref/>` markers, exposing a lossless expansion path via an injected tool.
 
-Status: **v0.14.0-alpha** — proof of concept. Measured 6.64× compression on a single 6.8 KB content block in smoke test. Known sharp edges below.
+Status: **v0.14.1** — request rewriting + transparent expand-tool resolution. Measured 6.64× wire-bandwidth compression on a single 6.8 KB content block; expand-tool round-trips are invisible to the client. Streaming + tool_use-block content + image content are still v0.14.2+ work.
 
 ## What it does
 
@@ -52,13 +52,14 @@ Smoke test against a mock upstream:
 
 Real-world LLM-token savings depend on how often the LLM resists calling `omc_proxy_expand_ref`. The tool's description tells the LLM to only expand when the preview isn't enough; in practice this should hold ~70-90% of the time on long contexts where most prior turns aren't load-bearing for the current response.
 
-## Known limitations (v0.14.0-alpha)
+## Known limitations (v0.14.1)
 
-1. **No streaming.** Requests with `"stream": true` pass through unchanged (no SSE rewriting yet).
-2. **The injected `omc_proxy_expand_ref` tool is not yet served by the proxy.** If the LLM actually emits a `tool_use` block calling it, the response flows back through Claude Code, which doesn't know the tool and will error. This means: in this alpha, the proxy works best in a "fire and forget" mode where the LLM responds without expanding markers. A v0.14.1 follow-up will catch tool_use for `omc_proxy_expand_ref` in the response stream, execute it locally from the cache, and inject the tool_result before returning to the client.
-3. **No image / tool_use / citation block rewriting.** Only `text` blocks and the `content` field of `tool_result` blocks (string or text-array form) are rewritten.
-4. **Response is forwarded unmodified.** Cache only fills on the request side, so the savings kick in on subsequent turns where prior big content reappears in conversation history.
+1. **No streaming.** Requests with `"stream": true` pass through unchanged (no SSE rewriting yet — v0.14.2 work).
+2. **Mixed tool_use passes through unchanged.** When the LLM emits the expand call alongside another tool call in the same response, the proxy doesn't intercept — it forwards the full response to the client (which will see the unknown expand tool). The auto-resolution only triggers when expand is the sole tool_use.
+3. **No image / `tool_use`-block / citation block rewriting.** Only `text` blocks and the `content` field of `tool_result` blocks (string or text-array form) are rewritten.
+4. **Response body is not cached for next-turn rewriting.** Cache only fills on the request side, so the savings kick in on subsequent turns where prior big content reappears in conversation history. A v0.15 follow-up will also index large assistant `text` blocks.
 5. **No batching API support.** `/v1/messages/batches` falls through to the generic passthrough.
+6. **Expand-loop bound at 8 rounds.** If the LLM keeps requesting expansion, the proxy gives up and returns 502 — protects against runaway tool-loop costs.
 
 ## Threat model
 
@@ -84,7 +85,8 @@ The proxy is a thin axum HTTP server. State lives in the existing OMC MemoryStor
 
 ## Roadmap
 
-- **v0.14.1**: catch `omc_proxy_expand_ref` tool_use in responses, execute locally
+- ~~**v0.14.1**: catch `omc_proxy_expand_ref` tool_use in responses, execute locally~~ ✅ shipped
 - **v0.14.2**: streaming SSE response rewriting
-- **v0.15.0**: tool_use / citation / image content support; batching API
+- **v0.14.3**: handle mixed tool_use (expand + other in same assistant turn)
+- **v0.15.0**: tool_use / citation / image content support; batching API; response-side caching for next-turn rewrites
 - **v0.16.0**: cache namespace per-conversation (use `x-conversation-id` or similar) so concurrent sessions don't collide

omnimcode-apiproxy/src/main.rs

@@ -125,7 +125,11 @@ async fn main() -> Result<()> {
     Ok(())
 }
 
-/// Rewrite-and-forward the /v1/messages POST.
+/// Rewrite-and-forward the /v1/messages POST. After receiving the upstream
+/// response, if the assistant emitted a sole tool_use for
+/// `omc_proxy_expand_ref`, the proxy resolves it locally from the cache and
+/// issues a follow-up upstream request — the client never sees the
+/// expand-tool round-trip. Mixed tool_use (expand + other) passes through.
 async fn handle_messages(State(state): State<AppState>, req: Request) -> Response {
     let (parts, body) = req.into_parts();
     let body_bytes = match axum::body::to_bytes(body, usize::MAX).await {
@@ -134,20 +138,17 @@ async fn handle_messages(State(state): State<AppState>, req: Request) -> Respons
             &format!("read request body: {}", e)),
     };
 
-    // Streaming requests bypass the rewriter for now — handling SSE responses
-    // requires per-event chunked rewriting which is v0.14.1+ work.
     let is_streaming = is_streaming_request(&body_bytes);
+    if is_streaming {
+        debug!("streaming request — bypassing rewriter (v0.14.2 work)");
+        return forward_to_upstream(&state, &parts.headers, body_bytes).await;
+    }
 
-    let rewritten = if is_streaming {
-        debug!("streaming request — passing through without rewrite");
-        body_bytes.clone()
-    } else {
-        match rewrite_request_body(&body_bytes, &state) {
-            Ok(b) => b,
-            Err(e) => {
-                warn!("rewrite failed, passing original through: {}", e);
-                body_bytes.clone()
-            }
+    let rewritten = match rewrite_request_body(&body_bytes, &state) {
+        Ok(b) => b,
+        Err(e) => {
+            warn!("rewrite failed, passing original through: {}", e);
+            body_bytes.clone()
         }
     };
 
@@ -157,7 +158,140 @@ async fn handle_messages(State(state): State<AppState>, req: Request) -> Respons
             body_bytes.len(), rewritten.len(), -saved);
     }
 
-    forward_to_upstream(&state, &parts.headers, rewritten).await
+    handle_with_expand_loop(&state, &parts.headers, rewritten).await
+}
+
+/// Upstream call + expand-tool auto-resolution loop. If the upstream's
+/// response contains a sole `tool_use` for `omc_proxy_expand_ref`, look
+/// up the hash in the cache, build a follow-up request with the
+/// tool_result synthetically appended, and re-call upstream. Bounded to
+/// MAX_EXPAND_ROUNDS to prevent runaway loops if the LLM keeps asking
+/// to expand.
+async fn handle_with_expand_loop(
+    state: &AppState, headers: &HeaderMap, initial_body: Bytes,
+) -> Response {
+    const MAX_EXPAND_ROUNDS: usize = 8;
+    let mut current_body = initial_body;
+    for round in 0..MAX_EXPAND_ROUNDS {
+        // Forward to upstream
+        let url = format!("{}/v1/messages",
+            state.upstream.trim_end_matches('/'));
+        let mut req = state.http.post(&url).body(current_body.to_vec());
+        for (k, v) in headers.iter() {
+            if k != "host" && k != "content-length" { req = req.header(k, v); }
+        }
+        let upstream_resp = match req.send().await {
+            Ok(r) => r,
+            Err(e) => return error_response(StatusCode::BAD_GATEWAY,
+                &format!("upstream: {}", e)),
+        };
+        let status = upstream_resp.status();
+        let resp_headers = upstream_resp.headers().clone();
+        let resp_body = match upstream_resp.bytes().await {
+            Ok(b) => b,
+            Err(e) => return error_response(StatusCode::BAD_GATEWAY,
+                &format!("read upstream: {}", e)),
+        };
+        // Only intercept successful, parseable responses
+        if !status.is_success() {
+            return rebuild_response(status, &resp_headers, resp_body);
+        }
+        let resp_json: Value = match serde_json::from_slice(&resp_body) {
+            Ok(v) => v,
+            Err(_) => return rebuild_response(status, &resp_headers, resp_body),
+        };
+
+        // Look for an exclusive expand tool_use
+        let expand_calls = collect_sole_expand_tool_uses(&resp_json);
+        if expand_calls.is_empty() {
+            return rebuild_response(status, &resp_headers, resp_body);
+        }
+        info!("round {}: auto-resolving {} expand tool_use(s)",
+            round + 1, expand_calls.len());
+
+        // Build follow-up request: previous messages + assistant response
+        // (rewritten through marker logic) + new user turn with tool_result
+        let mut next_req: Value = match serde_json::from_slice(&current_body) {
+            Ok(v) => v,
+            Err(_) => return rebuild_response(status, &resp_headers, resp_body),
+        };
+        let messages = next_req.get_mut("messages")
+            .and_then(Value::as_array_mut);
+        let Some(messages) = messages else {
+            return rebuild_response(status, &resp_headers, resp_body);
+        };
+        // Append the assistant turn (the upstream's response) verbatim
+        if let Some(asst_content) = resp_json.get("content").cloned() {
+            messages.push(json!({"role": "assistant", "content": asst_content}));
+        }
+        // Append a user turn with one tool_result per expand call
+        let mut tool_results: Vec<Value> = Vec::new();
+        for (tool_use_id, hash_str) in &expand_calls {
+            let body_text = lookup_expand(&hash_str, &state).unwrap_or_else(|e|
+                format!("[apiproxy: expand cache miss for {}: {}]", hash_str, e));
+            tool_results.push(json!({
+                "type": "tool_result",
+                "tool_use_id": tool_use_id,
+                "content": body_text,
+            }));
+        }
+        messages.push(json!({"role": "user", "content": tool_results}));
+
+        current_body = Bytes::from(serde_json::to_vec(&next_req).unwrap());
+    }
+    warn!("expand loop exceeded {} rounds, returning error", MAX_EXPAND_ROUNDS);
+    error_response(StatusCode::BAD_GATEWAY,
+        "apiproxy: expand loop limit exceeded")
+}
+
+/// If the response's `content` array contains exactly one tool_use AND it
+/// is for `omc_proxy_expand_ref`, return its (id, hash_str). Returning
+/// multiple results means there were multiple expand calls in a row, which
+/// also auto-resolves. Returns empty Vec for mixed tool_use (skip
+/// interception, let client handle) or no tool_use at all.
+fn collect_sole_expand_tool_uses(resp: &Value) -> Vec<(String, String)> {
+    let Some(content) = resp.get("content").and_then(Value::as_array) else {
+        return vec![];
+    };
+    let mut expand = Vec::new();
+    let mut has_other_tool_use = false;
+    for block in content {
+        if block.get("type").and_then(Value::as_str) == Some("tool_use") {
+            let name = block.get("name").and_then(Value::as_str).unwrap_or("");
+            if name == EXPAND_TOOL_NAME {
+                let id = block.get("id").and_then(Value::as_str)
+                    .unwrap_or("").to_string();
+                let hash = block.get("input")
+                    .and_then(|i| i.get("hash_str"))
+                    .and_then(Value::as_str).unwrap_or("").to_string();
+                if !id.is_empty() && !hash.is_empty() {
+                    expand.push((id, hash));
+                }
+            } else {
+                has_other_tool_use = true;
+            }
+        }
+    }
+    if has_other_tool_use { vec![] } else { expand }
+}
+
+fn lookup_expand(hash_str: &str, state: &AppState) -> Result<String> {
+    let hash: i64 = hash_str.parse()
+        .map_err(|e| anyhow::anyhow!("hash_str parse: {}", e))?;
+    let body = state.store.recall(Some(PROXY_CACHE_NAMESPACE), hash)
+        .map_err(anyhow::Error::msg)?
+        .ok_or_else(|| anyhow::anyhow!("not in cache"))?;
+    Ok(body)
+}
+
+fn rebuild_response(status: StatusCode, headers: &HeaderMap, body: Bytes) -> Response {
+    let mut resp = Response::builder().status(status);
+    for (k, v) in headers.iter() {
+        if k != "transfer-encoding" && k != "connection" && k != "content-length" {
+            resp = resp.header(k, v);
+        }
+    }
+    resp.body(axum::body::Body::from(body)).unwrap()
 }
 
 /// Forward anything else (model list, batches, etc.) unmodified.
@@ -178,44 +312,43 @@ async fn passthrough(State(state): State<AppState>, req: Request) -> Response {
         }
     }
     match req.send().await {
-        Ok(r) => relay(r).await,
-        Err(e) => error_response(StatusCode::BAD_GATEWAY, &format!("upstream: {}", e)),
+        Ok(r) => {
+            let status = r.status();
+            let h = r.headers().clone();
+            match r.bytes().await {
+                Ok(b) => rebuild_response(status, &h, b),
+                Err(e) => error_response(StatusCode::BAD_GATEWAY,
+                    &format!("read upstream: {}", e)),
+            }
+        }
+        Err(e) => error_response(StatusCode::BAD_GATEWAY,
+            &format!("upstream: {}", e)),
     }
 }
 
+/// Used by the streaming-passthrough path in handle_messages and by the
+/// catch-all passthrough route. Bytes-in, bytes-out, no rewriting.
 async fn forward_to_upstream(
     state: &AppState, headers: &HeaderMap, body: Bytes,
 ) -> Response {
     let url = format!("{}/v1/messages", state.upstream.trim_end_matches('/'));
     let mut req = state.http.post(&url).body(body.to_vec());
     for (k, v) in headers.iter() {
-        if k != "host" && k != "content-length" {
-            req = req.header(k, v);
-        }
+        if k != "host" && k != "content-length" { req = req.header(k, v); }
     }
     match req.send().await {
-        Ok(r) => relay(r).await,
-        Err(e) => error_response(StatusCode::BAD_GATEWAY, &format!("upstream: {}", e)),
-    }
-}
-
-async fn relay(upstream: reqwest::Response) -> Response {
-    let status = upstream.status();
-    let headers = upstream.headers().clone();
-    let body = match upstream.bytes().await {
-        Ok(b) => b,
-        Err(e) => return error_response(StatusCode::BAD_GATEWAY,
-            &format!("read upstream: {}", e)),
-    };
-    let mut resp = Response::builder().status(status);
-    for (k, v) in headers.iter() {
-        // Skip hop-by-hop headers; reqwest already drops Transfer-Encoding
-        // but Content-Length needs to track our (possibly rewritten) body.
-        if k != "transfer-encoding" && k != "connection" {
-            resp = resp.header(k, v);
+        Ok(r) => {
+            let status = r.status();
+            let h = r.headers().clone();
+            match r.bytes().await {
+                Ok(b) => rebuild_response(status, &h, b),
+                Err(e) => error_response(StatusCode::BAD_GATEWAY,
+                    &format!("read upstream: {}", e)),
+            }
         }
+        Err(e) => error_response(StatusCode::BAD_GATEWAY,
+            &format!("upstream: {}", e)),
     }
-    resp.body(axum::body::Body::from(body)).unwrap()
 }
 
 fn error_response(code: StatusCode, msg: &str) -> Response {