Skip to content

Commit 6e0ab58

Browse files
Osher-ElhadadOsher-Elhadad
authored andcommitted
Fix: Address context-guru plugin review feedback (#660)
Follow-up to the merged context-guru plugin (#660), addressing the non-blocking review comments from @mrsabath and @huang195. Plugin (authlib/plugins/contextguru): - OnRequest now skips compaction when there is no session identity instead of falling back to an empty ("") engine-store key, which would bleed per-session compaction state across unrelated callers. - Configure fails loud when exactly one of model.base_url / model.model is set (previously it silently degraded to deterministic — a typo'd or empty ${VAR} expansion would disable the cheap model with no error). - Add unit tests for the sentinel-header recursion guard (the one safety-critical branch) and the no-session skip. Build footprint: - Make context-guru opt-IN (//go:build include_plugin_contextguru) rather than opt-out. Its embedded engine pulls a large transitive set (bifrost/core, tiktoken-go, tree-sitter grammars, starlark), so the default authbridge-proxy/-envoy binaries stay lean (~16 MB smaller) and only pay for it when built with -tags include_plugin_contextguru. Documented in authbridge/CLAUDE.md as the deliberate exception to the exclude_plugin_* convention. Demo (demos/context-guru): - Remove the hardcoded internal IBM VPC endpoint from run.sh and agent.yaml; require CG_MODEL_BASE (fail-loud like CG_MODEL_KEY) and wire it into the deployment via kubectl set env, so no real endpoint lands in the repo. agent.yaml ships a neutral placeholder. - Harden the drive() poll loop: break on terminal pod phases (Failed/Unknown) with a readable error + describe/logs instead of piping crashed-pod logs into json.loads. - Add modest resources requests/limits to both demo containers so the manifest doubles as a usable template. - run.sh builds with -tags include_plugin_contextguru for the opt-in plugin. Supply-chain / hardening: - Digest-pin the envoy:v1.37.1 builder stage in the envoy Dockerfile (parity with the other stages). - Digest-pin alpine:3.20 in run.sh's generated demo Dockerfile. Assisted-By: Claude (Anthropic AI) <noreply@anthropic.com> Signed-off-by: Osher-Elhadad <Osher.Elhadad@ibm.com>
1 parent 468386d commit 6e0ab58

9 files changed

Lines changed: 137 additions & 24 deletions

File tree

authbridge/CLAUDE.md

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,10 @@ binaries with shared auth logic in `authlib/`:
1414
a2a-parser, mcp-parser, inference-parser, opa, sparc, ibac, token-broker).
1515
**Every** plugin is excludable via `-tags exclude_plugin_<name>` — one
1616
`plugins_<name>.go` file per plugin, gated by `//go:build !exclude_plugin_<name>`;
17-
`main.go` imports no plugin package directly.
17+
`main.go` imports no plugin package directly. **Exception:** `context-guru` is
18+
**opt-IN** (`//go:build include_plugin_contextguru`, not compiled by default),
19+
because its embedded engine pulls a large transitive dependency set; build with
20+
`-tags include_plugin_contextguru` to link it in.
1821
- `cmd/authbridge-envoy/` — envoy-sidecar mode. ext_proc gRPC server hooked
1922
into Envoy. Full plugin set. The `exclude_plugin_ibac` tag applies here too.
2023
- `authbridge-lite` (**image, not a separate binary**) — `cmd/authbridge-proxy`

authbridge/authlib/plugins/contextguru/plugin.go

Lines changed: 23 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -199,7 +199,16 @@ func (p *ContextGuru) Configure(raw json.RawMessage) error {
199199
if m.MaxTokens > 0 {
200200
p.modelMax = m.MaxTokens
201201
}
202-
if m.BaseURL != "" && m.Model != "" {
202+
switch {
203+
case m.BaseURL == "" && m.Model == "":
204+
// Empty model block → no static model; LLM components degrade to
205+
// deterministic. (Common when the config is present but the operator
206+
// hasn't wired an endpoint yet.)
207+
case m.BaseURL == "" || m.Model == "":
208+
// Exactly one set is almost always a typo/empty ${VAR} expansion that
209+
// would silently disable the cheap model. Fail loud instead.
210+
return fmt.Errorf("context-guru model config: base_url and model must both be set (got base_url=%q model=%q)", m.BaseURL, m.Model)
211+
default:
203212
p.static = cgModel{
204213
c: llmclient.New(llmclient.Options{
205214
Endpoint: m.BaseURL,
@@ -227,19 +236,27 @@ func (p *ContextGuru) OnRequest(ctx context.Context, pctx *pipeline.Context) pip
227236
if p.pipe == nil || len(pctx.Body) == 0 || !p.gated(pctx.Path) {
228237
return cont
229238
}
239+
// The engine Store is keyed by session for per-caller isolation. With no
240+
// session identity every such request would share one empty key and bleed
241+
// compaction state across unrelated callers, so skip compaction entirely.
242+
sid := sessionID(pctx)
243+
if sid == "" {
244+
slog.Debug("context-guru skipped request with no session identity", "path", pctx.Path)
245+
return cont
246+
}
230247
provider := providerFor(pctx.Path)
231248
models := cgcomponents.ModelSpec{
232249
Static: p.static,
233250
Incoming: p.incomingModel(pctx, provider),
234251
}
235252
before := len(pctx.Body)
236-
out, changed := apply.BodyWithModel(ctx, p.pipe, p.store, provider, pctx.Body, sessionID(pctx), false, models)
253+
out, changed := apply.BodyWithModel(ctx, p.pipe, p.store, provider, pctx.Body, sid, false, models)
237254
if changed && len(out) > 0 {
238255
pctx.SetBody(out)
239256
// Per-request byte-level view (the engine's token-level view is logged by
240257
// logEmitter.Run). The framework also emits a body-mutation session event.
241258
slog.Info("context-guru rewrote request body",
242-
"provider", provider, "path", pctx.Path, "session", sessionID(pctx),
259+
"provider", provider, "path", pctx.Path, "session", sid,
243260
"bytesBefore", before, "bytesAfter", len(out), "pctSaved", pct(before, len(out)))
244261
}
245262
return cont
@@ -316,7 +333,9 @@ func modelFromBody(body []byte) string {
316333
}
317334

318335
// sessionID keys per-session engine state to the conversation: the A2A contextId
319-
// when present, else the caller subject for multi-tenant isolation.
336+
// when present, else the caller subject for multi-tenant isolation. Returns ""
337+
// when neither is available; OnRequest treats that as "skip compaction" rather
338+
// than sharing one store key across unrelated callers.
320339
func sessionID(pctx *pipeline.Context) string {
321340
if pctx.Session != nil && pctx.Session.ID != "" {
322341
return pctx.Session.ID

authbridge/authlib/plugins/contextguru/plugin_test.go

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -117,6 +117,48 @@ func TestOnRequest_CompactsLargeToolOutput(t *testing.T) {
117117
}
118118
}
119119

120+
// TestOnRequest_SentinelHeader_ShortCircuits covers the one safety-critical
121+
// branch: the plugin's own extract:code LLM call carries the sentinel, and if it
122+
// ever transits this forward proxy OnRequest must bail before touching the body —
123+
// otherwise the compaction call recurses into itself.
124+
func TestOnRequest_SentinelHeader_ShortCircuits(t *testing.T) {
125+
p := configured(t, collapseEngine)
126+
body := chatBody(bigToolOutput())
127+
pctx := inferencePctx("/v1/chat/completions", body)
128+
pctx.Headers.Set(sentinelHeader, "1")
129+
if act := invokeReq(p, pctx); act.Type != pipeline.Continue {
130+
t.Fatalf("expected Continue, got %v", act.Type)
131+
}
132+
if pctx.BodyMutated() {
133+
t.Fatal("sentinel header must short-circuit before any compaction")
134+
}
135+
}
136+
137+
// TestOnRequest_NoSession_Skips verifies that a request with neither session nor
138+
// identity is passed through untouched, so sessionless callers can't share one
139+
// empty engine-store key.
140+
func TestOnRequest_NoSession_Skips(t *testing.T) {
141+
p := configured(t, collapseEngine)
142+
body := chatBody(bigToolOutput())
143+
pctx := inferencePctx("/v1/chat/completions", body)
144+
pctx.Session = nil // and Identity is nil → sessionID == ""
145+
if act := invokeReq(p, pctx); act.Type != pipeline.Continue {
146+
t.Fatalf("expected Continue, got %v", act.Type)
147+
}
148+
if pctx.BodyMutated() {
149+
t.Fatal("no session identity must skip compaction (store-key isolation)")
150+
}
151+
}
152+
153+
func TestConfigure_RejectsPartialModel(t *testing.T) {
154+
if err := New().Configure(json.RawMessage(`{"model":{"base_url":"http://x"}}`)); err == nil {
155+
t.Fatal("expected error when model.model is unset but base_url is set")
156+
}
157+
if err := New().Configure(json.RawMessage(`{"model":{"model":"gpt-4o-mini"}}`)); err == nil {
158+
t.Fatal("expected error when model.base_url is unset but model is set")
159+
}
160+
}
161+
120162
func TestOnRequest_EmptyBody_PassThrough(t *testing.T) {
121163
p := configured(t, collapseEngine)
122164
pctx := inferencePctx("/v1/chat/completions", nil)

authbridge/cmd/authbridge-envoy/Dockerfile

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ ENV GOWORK=off
2727
RUN cd cmd/authbridge-envoy && CGO_ENABLED=0 GOOS=linux go build -tags "${GO_BUILD_TAGS}" -ldflags="-s -w" -o /authbridge-envoy .
2828

2929
# Stage 2: Get Envoy binary. Already stripped upstream; nothing to do.
30-
FROM docker.io/envoyproxy/envoy:v1.37.1 AS envoy-source
30+
FROM docker.io/envoyproxy/envoy:v1.37.1@sha256:29496a88fba9c4c9cdef4afe8fec70f536c5ba111b1c2bddbc5436b091ceca33 AS envoy-source
3131

3232
# Stage 3: Runtime — UBI9-micro (glibc-compatible, ~24 MB, has bash)
3333
FROM registry.access.redhat.com/ubi9/ubi-micro@sha256:2173487b3b72b1a7b11edc908e9bbf1726f9df46a4f78fd6d19a2bab0a701f38
Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,10 @@
1-
//go:build !exclude_plugin_contextguru
1+
//go:build include_plugin_contextguru
22

3+
// context-guru is opt-IN (positive build tag), unlike the other plugins which are
4+
// compiled in by default and dropped via exclude_plugin_*. Its embedded engine
5+
// pulls a large transitive set (bifrost/core, tiktoken-go, tree-sitter grammars,
6+
// starlark), so it is kept out of the default authbridge-proxy/-envoy binaries and
7+
// linked only when built with -tags include_plugin_contextguru.
38
package main
49

510
import _ "github.com/kagenti/kagenti-extensions/authbridge/authlib/plugins/contextguru"
Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,10 @@
1-
//go:build !exclude_plugin_contextguru
1+
//go:build include_plugin_contextguru
22

3+
// context-guru is opt-IN (positive build tag), unlike the other plugins which are
4+
// compiled in by default and dropped via exclude_plugin_*. Its embedded engine
5+
// pulls a large transitive set (bifrost/core, tiktoken-go, tree-sitter grammars,
6+
// starlark), so it is kept out of the default authbridge-proxy/-envoy binaries and
7+
// linked only when built with -tags include_plugin_contextguru.
38
package main
49

510
import _ "github.com/kagenti/kagenti-extensions/authbridge/authlib/plugins/contextguru"

authbridge/demos/context-guru/README.md

Lines changed: 22 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -91,19 +91,34 @@ And the observed answers:
9191
## Run it
9292

9393
```bash
94-
export CG_MODEL_KEY=<api key for the extract-code model> # e.g. an OpenAI-compatible key
94+
export CG_MODEL_KEY=<api key for the extract-code model> # e.g. an OpenAI-compatible key
95+
export CG_MODEL_BASE=https://api.openai.com # any OpenAI-wire endpoint
9596
./run.sh all # setup + drive off, observe, enforce and print the comparison
9697
# or step by step:
9798
./run.sh setup
9899
./run.sh drive enforce # (or observe / off)
99100
```
100101

101-
`run.sh setup` builds the `authbridge-proxy` image **with the context-guru plugin**,
102-
loads it + the enhanced `finance-mcp` into the `kagenti` kind cluster, creates a
103-
12,288-token-window Ollama model (`llama3.2-ctx12k`) so the raw request truncates,
104-
and deploys the agent + sidecar. `run.sh drive <mode>` flips `on_error`, restarts,
105-
drives the audit, and prints the agent's answer + the byte/token gain from the
106-
sidecar session API (`:9094`).
102+
`run.sh setup` builds the `authbridge-proxy` image **with the context-guru plugin**
103+
(`-tags include_plugin_contextguru` — see *Build integration* below), loads it + the
104+
enhanced `finance-mcp` into the `kagenti` kind cluster, creates a 12,288-token-window
105+
Ollama model (`llama3.2-ctx12k`) so the raw request truncates, and deploys the agent
106+
+ sidecar. `run.sh drive <mode>` flips `on_error`, restarts, drives the audit, and
107+
prints the agent's answer + the byte/token gain from the sidecar session API (`:9094`).
108+
109+
## Build integration
110+
111+
context-guru is **opt-in**: unlike the other plugins (compiled in by default,
112+
dropped via `-tags exclude_plugin_*`), it is linked only when the binary is built
113+
with `-tags include_plugin_contextguru`. Its embedded engine pulls a large
114+
transitive dependency set (bifrost/core, tiktoken-go, tree-sitter grammars,
115+
starlark), so the default `authbridge-proxy`/`authbridge-envoy` binaries stay lean
116+
and a deployment that doesn't want compaction never pays for it.
117+
118+
```bash
119+
cd authbridge && podman build -f cmd/authbridge-proxy/Dockerfile \
120+
--build-arg GO_BUILD_TAGS=include_plugin_contextguru -t authbridge-cg:latest .
121+
```
107122

108123
## Configuration reference
109124

authbridge/demos/context-guru/k8s/agent.yaml

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -39,17 +39,25 @@ spec:
3939
- { name: FINANCE_MCP_URL, value: "http://finance-mcp.team1.svc.cluster.local:8888" }
4040
- { name: AGENT_PUBLIC_URL, value: "http://cg-finance-agent.team1.svc.cluster.local:8080/" }
4141
volumeMounts: [{ mountPath: /shared, name: shared-data }]
42+
resources:
43+
requests: { cpu: "100m", memory: "128Mi" }
44+
limits: { cpu: "1", memory: "512Mi" }
4245
- name: authbridge-proxy
4346
image: authbridge-cg:latest # authbridge-proxy built with the context-guru plugin (see run.sh)
4447
imagePullPolicy: IfNotPresent
4548
args: ["--config", "/etc/authbridge/config.yaml"]
4649
env:
47-
- { name: CG_MODEL_BASE, value: "https://ete-litellm.ai-models.vpc-int.res.ibm.com" }
48-
- { name: CG_MODEL_NAME, value: "claude-haiku-4-5" } # capable, cheap extract-code model
50+
# Neutral placeholders — run.sh overrides these via `kubectl set env`
51+
# from CG_MODEL_BASE / CG_MODEL_NAME so no real endpoint lands in the repo.
52+
- { name: CG_MODEL_BASE, value: "https://api.openai.com" }
53+
- { name: CG_MODEL_NAME, value: "gpt-4o-mini" } # capable, cheap extract-code model
4954
- name: CG_MODEL_KEY
5055
valueFrom: { secretKeyRef: { name: cg-model-key, key: key } }
5156
ports: [{ containerPort: 8080 }, { containerPort: 8081 }, { containerPort: 9094 }]
5257
volumeMounts: [{ mountPath: /etc/authbridge, name: ab-config }]
58+
resources:
59+
requests: { cpu: "50m", memory: "64Mi" }
60+
limits: { cpu: "500m", memory: "256Mi" }
5361
volumes:
5462
- { name: shared-data, emptyDir: {} }
5563
- { name: ab-config, configMap: { name: cg-authbridge-config } }

authbridge/demos/context-guru/run.sh

Lines changed: 23 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -21,20 +21,20 @@ HERE="$(cd "$(dirname "$0")" && pwd)"
2121
AB="$(cd "$HERE/../.." && pwd)" # authbridge/
2222
FS="$AB/demos/finance-sparc"
2323
OLLAMA_HOST_URL="${OLLAMA_HOST_URL:-http://127.0.0.1:11434}" # host Ollama, as seen from THIS machine
24-
: "${CG_MODEL_BASE:=https://ete-litellm.ai-models.vpc-int.res.ibm.com}"
25-
: "${CG_MODEL_NAME:=claude-haiku-4-5}"
24+
: "${CG_MODEL_NAME:=gpt-4o-mini}" # any capable, cheap OpenAI-wire model
2625
CTX_MODEL="llama3.2-ctx12k"
2726

2827
setup() {
2928
: "${CG_MODEL_KEY:?set CG_MODEL_KEY to the extract-code model API key}"
29+
: "${CG_MODEL_BASE:?set CG_MODEL_BASE to an OpenAI-compatible endpoint for the extract-code model, e.g. https://api.openai.com}"
3030
local arch; arch=$(docker exec kagenti-control-plane uname -m | sed 's/aarch64/arm64/;s/x86_64/amd64/')
3131

32-
echo "==> build authbridge-proxy WITH the context-guru plugin ($arch, pure-Go) + image"
32+
echo "==> build authbridge-proxy WITH the context-guru plugin ($arch, pure-Go, -tags include_plugin_contextguru) + image"
3333
( cd "$AB" && GOTOOLCHAIN=auto GOOS=linux GOARCH="$arch" CGO_ENABLED=0 \
34-
go build -ldflags="-s -w" -o /tmp/authbridge-proxy-cg ./cmd/authbridge-proxy )
34+
go build -tags include_plugin_contextguru -ldflags="-s -w" -o /tmp/authbridge-proxy-cg ./cmd/authbridge-proxy )
3535
local d=/tmp/cg-img; mkdir -p "$d"; cp /tmp/authbridge-proxy-cg "$d/authbridge-proxy"
3636
cat > "$d/Dockerfile" <<'DOCKER'
37-
FROM alpine:3.20
37+
FROM alpine:3.20@sha256:beefdbd8a1da6d2915566fde36db9db0b524eb737fc57cd1367effd16dc0d06d
3838
RUN apk add --no-cache ca-certificates
3939
COPY authbridge-proxy /usr/local/bin/authbridge-proxy
4040
ENTRYPOINT ["/usr/local/bin/authbridge-proxy"]
@@ -59,6 +59,10 @@ DOCKER
5959
--dry-run=client -o yaml | kubectl apply -f -
6060
kubectl apply -f "$HERE/k8s/agent.yaml"
6161
kubectl -n "$NS" set env deploy/cg-finance-agent OLLAMA_MODEL="$CTX_MODEL" >/dev/null
62+
# Point the extract-code model at the operator-provided endpoint (agent.yaml
63+
# ships only a neutral placeholder; the real endpoint never lands in the repo).
64+
kubectl -n "$NS" set env deploy/cg-finance-agent -c authbridge-proxy \
65+
CG_MODEL_BASE="$CG_MODEL_BASE" CG_MODEL_NAME="$CG_MODEL_NAME" >/dev/null
6266
kubectl -n "$NS" rollout status deploy/cg-finance-agent --timeout=120s
6367
}
6468

@@ -77,8 +81,20 @@ JSON
7781
kubectl -n "$NS" delete pod cgdrive --ignore-not-found >/dev/null; sleep 2
7882
kubectl -n "$NS" run cgdrive --restart=Never --image=curlimages/curl:8.10.1 \
7983
--overrides='{"spec":{"containers":[{"name":"c","image":"curlimages/curl:8.10.1","command":["sh","-c","curl -m 300 -s -X POST http://cg-finance-agent.team1.svc.cluster.local:8080/ -H '"'"'Content-Type: application/json'"'"' --data @/data/drive.json"],"volumeMounts":[{"name":"d","mountPath":"/data"}]}],"volumes":[{"name":"d","configMap":{"name":"cg-drive"}}]}}' >/dev/null
80-
while [ "$(kubectl -n "$NS" get pod cgdrive -o jsonpath='{.status.phase}')" = "Running" ] || \
81-
[ "$(kubectl -n "$NS" get pod cgdrive -o jsonpath='{.status.phase}')" = "Pending" ]; do sleep 5; done
84+
# Wait for the driver pod to finish. Break explicitly on terminal failure so a
85+
# crashed/unschedulable pod gives a readable error instead of piping non-JSON
86+
# logs into json.loads below.
87+
while :; do
88+
phase=$(kubectl -n "$NS" get pod cgdrive -o jsonpath='{.status.phase}' 2>/dev/null || echo Unknown)
89+
case "$phase" in
90+
Succeeded) break ;;
91+
Running|Pending) sleep 5 ;;
92+
*) echo "!! cgdrive pod ended in phase=$phase — the audit did not complete:"
93+
kubectl -n "$NS" describe pod cgdrive | tail -n 25
94+
kubectl -n "$NS" logs cgdrive 2>/dev/null || true
95+
exit 1 ;;
96+
esac
97+
done
8298
echo "--- agent answer ($mode) ---"
8399
kubectl -n "$NS" logs cgdrive | python3 -c 'import sys,json;b=sys.stdin.read().split("EXIT=")[0].strip();d=json.loads(b);r=d.get("result",{});m=(r.get("status")or{}).get("message")or{};p=m.get("parts") or (r.get("artifacts") or [{}])[0].get("parts",[]);print(" ".join(x.get("text","") for x in p)[:600])'
84100
echo "--- context-guru gain ($mode) ---"

0 commit comments

Comments
 (0)