Skip to content

feat: add config plugin support with additional header injections #2

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
schneidermr merged 3 commits into from
Mar 23, 2026
Merged

feat: add config plugin support with additional header injections #2

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
3b1f64b
feat: handle the additional header elements injected by the IDE plugins
schneidermr Mar 23, 2026
34c8453
feat: store unprocessed events separate from audit
schneidermr Mar 23, 2026
0083d47
docs: update the readme and the changelog
schneidermr Mar 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 25 additions & 3 deletions CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,10 +7,30 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0

## [Unreleased]

## [0.4.0] - 2026-03-23

### Added

- `ANTHROPIC_UPSTREAM_URL`, `OPENAI_UPSTREAM_URL`, `GEMINI_UPSTREAM_URL` env vars — override the default upstream target for each provider without changing detection logic; useful for routing through LiteLLM, Portkey, a regional endpoint, or any internal gateway; defaults remain `https://api.anthropic.com`, `https://api.openai.com`, `https://generativelanguage.googleapis.com`
- `NewRouterWithConfig(RouterConfig)` constructor on `Router` — accepts the three upstream URL overrides alongside the existing SAP AI Core config; `NewRouter` is preserved as a backwards-compatible wrapper
- **IDE plugin header support** — `X-Audit-User`, `X-Audit-Team`, `X-Audit-Project`, `X-Audit-Branch`, `X-Audit-Session-ID`, and `X-Audit-Agent` headers allow IDE companion plugins to inject per-request identity and context; header values override env-var defaults (`AUDIT_PROJECT`, `AUDIT_BRANCH`) for centrally-hosted deployments
- **Unprocessed events table** — responses the proxy cannot parse into structured audit records (non-chat endpoints like `/v1/models` and `/v1/count_tokens`, unknown providers, parse errors) are routed to `audit.unprocessed_events` instead of being silently dropped; stores raw body, HTTP method, path, status code, content type, and parse error message
- **Health endpoint** — `GET /healthz` returns `{"status":"ok","version":"..."}` for IDE plugin connectivity checks and load balancer probes; version set at build time via `-ldflags "-X main.version=..."`
- `user` and `team` columns on `audit_events` — `LowCardinality(String)`, populated from `X-Audit-User` and `X-Audit-Team` headers
- `UnprocessedAdder` interface and `UnprocessedBatcher` — mirrors `EventAdder`/`Batcher` pattern for the unprocessed events write path
- `ANTHROPIC_UPSTREAM_URL`, `OPENAI_UPSTREAM_URL`, `GEMINI_UPSTREAM_URL` env vars — override the default upstream target for each provider without changing detection logic; useful for routing through LiteLLM, Portkey, a regional endpoint, or any internal gateway
- `NewRouterWithConfig(RouterConfig)` constructor on `Router` — accepts upstream URL overrides alongside SAP AI Core config; `NewRouter` preserved as backwards-compatible wrapper
- Migration `005_add_user_team.sql` — adds `user` and `team` columns to existing installations
- Migration `006_unprocessed_events.sql` — creates the `unprocessed_events` table
- Cline agent detection — `detectAgent` now matches `cline` in the User-Agent header
- Claude Code VSCode extension detection — `detectAgent` now matches `claude-cli` User-Agent sent by the VSCode extension (in addition to `claude-code` from the CLI)
- Session ID priority chain: `X-Audit-Session-ID` > `X-Session-ID` > auto-generated UUID
- IDE plugins design document (`docs/ide-plugins-design.md`) — comprehensive design for VSCode and JetBrains companion plugins covering env var injection, config file schema, and header protocol
- New handler tests: `X-Audit-*` header overrides, session ID fallback, Cline detection, Claude VSCode extension detection, unprocessed event routing

### Changed

- `extractEvents` now returns `([]audit.AuditEvent, error)` instead of `[]audit.AuditEvent` — parse errors are surfaced to the caller for routing to the unprocessed table
- `NewHandler` signature extended with `unprocessed audit.UnprocessedAdder` parameter
- `requestMeta` struct extended with `user`, `team`, `method`, and `unprocessed` fields

## [0.3.0] - 2026-03-18

Expand Down Expand Up @@ -71,6 +91,8 @@ Initial release.
- ClickHouse migrations: `001_initial.sql` (full schema), `002_add_branch.sql`
- Test suite: parser unit tests, batcher tests, router tests, stream tap tests, handler integration tests

[Unreleased]: https://github.com/bitkaio/codesteward-audit-proxy/compare/v0.2.0...HEAD
[Unreleased]: https://github.com/bitkaio/codesteward-audit-proxy/compare/v0.4.0...HEAD
[0.4.0]: https://github.com/bitkaio/codesteward-audit-proxy/compare/v0.3.0...v0.4.0
[0.3.0]: https://github.com/bitkaio/codesteward-audit-proxy/compare/v0.2.0...v0.3.0
[0.2.0]: https://github.com/bitkaio/codesteward-audit-proxy/compare/v0.1.0...v0.2.0
[0.1.0]: https://github.com/bitkaio/codesteward-audit-proxy/releases/tag/v0.1.0
115 changes: 98 additions & 17 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,10 +22,13 @@ The proxy is fully transparent. It never buffers the response before forwarding,
## Features

- **Stream tap, never buffer** — uses `io.TeeReader` to forward tokens to the agent immediately while capturing a copy for audit asynchronously
- **Anthropic + OpenAI parsing** — extracts thinking blocks, text, and tool calls from both streaming (SSE) and non-streaming responses
- **Anthropic + OpenAI + SAP AI Core parsing** — extracts thinking blocks, text, and tool calls from both streaming (SSE) and non-streaming responses
- **Request capture with scrubbing** — records user-role messages in a structured `user_messages` column; configurable regexp scrubbing replaces sensitive content with `[REDACTED]` before storage
- **IDE plugin header support** — `X-Audit-User`, `X-Audit-Team`, `X-Audit-Project`, `X-Audit-Branch`, and `X-Audit-Session-ID` headers allow IDE companion plugins to inject per-request identity and context
- **Unprocessed event capture** — responses that cannot be parsed into structured audit records (non-chat endpoints, unknown providers, parse errors) are routed to a separate `unprocessed_events` table so no data is lost
- **Health endpoint** — `GET /healthz` returns JSON status and version for IDE plugin connectivity checks and load balancer probes
- **Batched ClickHouse writes** — accumulates events in memory and flushes on size threshold (default 100) or time interval (default 1s)
- **Multi-tenancy** — `AUDIT_PROJECT` and `AUDIT_BRANCH` tag every row so multiple repos and branches share one ClickHouse instance
- **Multi-tenancy** — `AUDIT_PROJECT` and `AUDIT_BRANCH` tag every row; `X-Audit-*` headers override env-var defaults for centrally-hosted deployments
- **OpenTelemetry traces** — one span per proxied request with `gen_ai.system`, session/turn IDs, latency, and status; flush spans per ClickHouse batch; W3C trace context propagated in both directions
- **Proxy chaining** — supports `UPSTREAM_PROXY` for corporate firewalls, Portkey, LiteLLM, and other gateway proxies
- **Structured JSON logging** — every request, batch flush, and error logged via `log/slog`
Expand Down Expand Up @@ -166,6 +169,8 @@ All configuration is via environment variables. No config files required.
| `ANTHROPIC_UPSTREAM_URL` | `https://api.anthropic.com` | Override Anthropic target (e.g. LiteLLM, Portkey, regional endpoint) |
| `OPENAI_UPSTREAM_URL` | `https://api.openai.com` | Override OpenAI target |
| `GEMINI_UPSTREAM_URL` | `https://generativelanguage.googleapis.com` | Override Gemini target |
| `SAP_AICORE_BASE_URL` | *(none)* | SAP AI Core API URL; enables SAP AI Core routing when set |
| `SAP_AICORE_AUTH_HOST` | `ml.hana.ondemand.com` | Host fragment for detecting SAP AI Core traffic |
| `UPSTREAM_PROXY` | *(none)* | Upstream proxy URL (overrides `HTTPS_PROXY`) |
| `LOG_LEVEL` | `info` | `debug` \| `info` \| `warn` \| `error` |
| `OTEL_EXPORTER_OTLP_ENDPOINT` | *(none)* | Activates OTel traces when set (e.g. `http://localhost:4318`) |
Expand Down Expand Up @@ -223,9 +228,38 @@ The proxy extracts W3C `traceparent`/`tracestate` headers from incoming agent re

---

## Health endpoint

`GET /healthz` returns a JSON response with the proxy status and build version:

```json
{"status": "ok", "version": "0.4.0"}
```

Useful for IDE plugin connectivity checks, load balancer probes, and deployment verification. The version is set at build time via `-ldflags "-X main.version=..."`.

---

## IDE plugin headers

When the proxy is hosted centrally (shared by a team), individual developers cannot set env vars on the proxy process. Instead, IDE companion plugins (VSCode, JetBrains) inject per-request identity via `X-Audit-*` headers.

| Header | Description | Override behaviour |
| --- | --- | --- |
| `X-Audit-User` | Developer identity (git email, username) | Stored as-is |
| `X-Audit-Team` | Team or org identifier | Stored as-is |
| `X-Audit-Project` | Repository / project name | Overrides `AUDIT_PROJECT` env var |
| `X-Audit-Branch` | Git branch | Overrides `AUDIT_BRANCH` env var |
| `X-Audit-Session-ID` | Session identifier | Takes priority over `X-Session-ID` and auto-generated UUID |
| `X-Audit-Agent` | Agent name override | Overrides User-Agent detection |

All `X-Audit-*` headers are stripped before forwarding to upstream APIs — they never reach the LLM provider.

---

## Multi-tenancy

`AUDIT_PROJECT` and `AUDIT_BRANCH` tag every ClickHouse row, allowing multiple repositories and branches to share one instance.
`AUDIT_PROJECT` and `AUDIT_BRANCH` tag every ClickHouse row, allowing multiple repositories and branches to share one instance. For centrally-hosted proxies, `X-Audit-Project` and `X-Audit-Branch` headers override the env-var defaults on a per-request basis.

```bash
export AUDIT_PROJECT=myorg/myrepo
Expand All @@ -237,10 +271,10 @@ go run ./cmd/proxy
Example query across tenants:

```sql
SELECT project, branch, agent, tool_name, count() AS calls
SELECT project, branch, user, agent, tool_name, count() AS calls
FROM audit.audit_events
WHERE toDate(ts) = today()
GROUP BY project, branch, agent, tool_name
GROUP BY project, branch, user, agent, tool_name
ORDER BY calls DESC;
```

Expand All @@ -255,6 +289,7 @@ The proxy detects the upstream from the `Host` header first, then the request pa
| `api.anthropic.com` or `/v1/messages` | `https://api.anthropic.com` |
| `api.openai.com` or `/v1/chat/` | `https://api.openai.com` |
| `generativelanguage.googleapis.com` or `/v1beta/` | `https://generativelanguage.googleapis.com` |
| SAP AI Core host (configurable) | SAP AI Core base URL |

---

Expand All @@ -278,6 +313,8 @@ go run ./cmd/proxy

## ClickHouse schema

### audit_events

One row per tool call. Responses with no tool calls produce a single row with `tool_name = ''`. Request-direction rows carry `user_messages` and have `direction = 'request'`.

```sql
Expand All @@ -293,52 +330,96 @@ CREATE TABLE audit.audit_events
thinking Array(String),
assistant_text Array(String),
tool_name String,
tool_input String, -- JSON-encoded tool input
tool_input String, -- JSON-encoded tool input
model LowCardinality(String),
raw String, -- full original body (scrubbed if patterns set)
request_captured UInt8, -- 0 when AUDIT_CAPTURE_REQUESTS=false
user_messages Array(String) -- extracted user-role text, scrubbed
raw String, -- full original body (scrubbed if patterns set)
resource_group String, -- SAP AI Core resource group
request_captured UInt8, -- 0 when AUDIT_CAPTURE_REQUESTS=false
user_messages Array(String), -- extracted user-role text, scrubbed
user LowCardinality(String), -- developer identity from X-Audit-User
team LowCardinality(String) -- team/org from X-Audit-Team
)
ENGINE = MergeTree()
PARTITION BY toYYYYMM(ts)
ORDER BY (project, session_id, ts);
```

Existing installations: apply migrations in order:
### unprocessed_events

Responses (and requests) that cannot be parsed into structured audit records are stored here. This includes non-chat endpoints (e.g. `/v1/models`, `/v1/count_tokens`), unknown providers, and parse errors.

```sql
CREATE TABLE audit.unprocessed_events
(
session_id String,
turn_id String,
ts DateTime64(3),
agent LowCardinality(String),
project String,
branch LowCardinality(String),
user LowCardinality(String),
team LowCardinality(String),
direction LowCardinality(String),
method LowCardinality(String),
path String,
status_code UInt16,
content_type LowCardinality(String),
raw String,
error String
)
ENGINE = MergeTree()
PARTITION BY toYYYYMM(ts)
ORDER BY (project, session_id, ts);
```

### Migrations

Existing installations: apply migrations in order. If using Docker Compose, migrations are applied automatically.

```bash
clickhouse-client --multiquery < migrations/002_add_branch.sql
clickhouse-client --multiquery < migrations/003_request_capture.sql
clickhouse-client --multiquery < migrations/004_sapaicore.sql
clickhouse-client --multiquery < migrations/005_add_user_team.sql
clickhouse-client --multiquery < migrations/006_unprocessed_events.sql
```

---

## Repository structure

```text
├── cmd/proxy/main.go Entry point, wiring, graceful shutdown
├── cmd/proxy/main.go Entry point, /healthz, wiring, graceful shutdown
├── internal/
│ ├── config/config.go Env-var config loading, git branch detection
│ ├── telemetry/otel.go OTel TracerProvider setup (no-op when unconfigured)
│ ├── audit/
│ │ ├── event.go AuditEvent struct + EventAdder interface
│ │ ├── batcher.go In-memory batcher (size + interval flush)
│ │ ├── event.go AuditEvent, UnprocessedEvent, EventAdder, UnprocessedAdder
│ │ ├── batcher.go In-memory batchers (size + interval flush)
│ │ ├── scrubber.go Scrubber interface, NopScrubber, PatternScrubber
│ │ └── clickhouse.go ClickHouse native-protocol writer
│ │ └── clickhouse.go ClickHouse native-protocol writer (audit + unprocessed)
│ ├── proxy/
│ │ ├── handler.go Reverse proxy handler, audit transport, OTel spans
│ │ ├── router.go Upstream detection and URL rewriting
│ │ ├── stream.go TeeReader stream tap
│ │ └── transport.go http.Transport with proxy chaining
│ └── parser/
│ ├── types.go Shared ToolCall type
│ ├── anthropic.go Anthropic message + SSE stream parser
│ ├── openai.go OpenAI chat completion + SSE stream parser
│ ├── sapaicore.go SAP AI Core response parser
│ ├── request.go Provider-agnostic request parser (user message extraction)
│ └── gemini.go Gemini stub (TODO)
├── docs/
│ └── ide-plugins-design.md VSCode + JetBrains companion plugin design
└── migrations/
├── 001_initial.sql Full schema for new installations
├── 002_add_branch.sql Add branch column to existing installations
└── 003_request_capture.sql Add request_captured + user_messages columns
├── 002_add_branch.sql Add branch column
├── 003_request_capture.sql Add request_captured + user_messages columns
├── 004_sapaicore.sql Add resource_group column
├── 005_add_user_team.sql Add user + team columns
├── 006_unprocessed_events.sql Create unprocessed_events table
└── migrate.sh Idempotent HTTP migration runner
```

---
Expand All @@ -349,7 +430,7 @@ clickhouse-client --multiquery < migrations/003_request_capture.sql
go test ./...
```

Tests cover: Anthropic and OpenAI parsing (full and streaming, including edge cases), batcher (size-threshold flush, ticker flush, drain on stop, non-blocking drop), scrubber (pattern redaction, multi-pattern, passthrough, invalid pattern error), request parser (string and array content, mixed conversations, scrubber application), router (host-based, path-based, header-based routing, URL rewriting), stream tap (byte fidelity, SSE detection, callback timing), and the handler end-to-end (status/body/header passthrough, internal header scrubbing, audit event emission, 502 on dead upstream).
Tests cover: Anthropic and OpenAI parsing (full and streaming, including edge cases), batcher (size-threshold flush, ticker flush, drain on stop, non-blocking drop), scrubber (pattern redaction, multi-pattern, passthrough, invalid pattern error), request parser (string and array content, mixed conversations, scrubber application), router (host-based, path-based, header-based routing, URL rewriting), stream tap (byte fidelity, SSE detection, callback timing), and the handler end-to-end (status/body/header passthrough, internal header scrubbing, audit event emission, X-Audit-* header overrides, unprocessed event routing, agent detection for all supported agents, 502 on dead upstream).

---

Expand Down
Binary file modified assets/codesteward-logo.png
View file
Open in desktop
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
24 changes: 21 additions & 3 deletions cmd/proxy/main.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package main

import (
"context"
"encoding/json"
"errors"
"log/slog"
"net/http"
Expand All @@ -17,6 +18,9 @@ import (
"llm-audit-proxy/internal/telemetry"
)

// version is set at build time via -ldflags "-X main.version=...".
var version = "dev"

func main() {
cfg, err := config.Load()
if err != nil {
Expand Down Expand Up @@ -67,8 +71,9 @@ func main() {
os.Exit(1)
}

// Batcher and transport.
// Batchers and transport.
batcher := audit.NewBatcher(writer, cfg.BatchSize, cfg.BatchInterval)
unprocessedBatcher := audit.NewUnprocessedBatcher(writer, cfg.BatchSize, cfg.BatchInterval)
transport := proxy.BuildTransport(cfg)

if cfg.SAPAICoreBaseURL == "" {
Expand All @@ -81,11 +86,23 @@ func main() {
SAPAICoreBaseURL: cfg.SAPAICoreBaseURL,
SAPAICoreAuthHost: cfg.SAPAICoreAuthHost,
})
handler := proxy.NewHandler(batcher, transport, cfg.AuditProject, cfg.AuditBranch, scrubber, cfg.CaptureRequests, router)
handler := proxy.NewHandler(batcher, transport, cfg.AuditProject, cfg.AuditBranch, scrubber, cfg.CaptureRequests, router, unprocessedBatcher)

// Top-level mux: /healthz is handled directly; everything else goes to
// the reverse proxy handler.
mux := http.NewServeMux()
mux.HandleFunc("GET /healthz", func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
json.NewEncoder(w).Encode(map[string]string{
"status": "ok",
"version": version,
})
})
mux.Handle("/", handler)

srv := &http.Server{
Addr: cfg.ProxyAddr,
Handler: handler,
Handler: mux,
ReadTimeout: 5 * time.Minute,
WriteTimeout: 5 * time.Minute,
IdleTimeout: 2 * time.Minute,
Expand Down Expand Up @@ -114,6 +131,7 @@ func main() {
}

batcher.Stop()
unprocessedBatcher.Stop()

if err := otelShutdown(shutdownCtx); err != nil {
slog.Error("otel shutdown error", "err", err)
Expand Down
Loading
Loading