pem: direct-query gRPC endpoint — stub + TDD contract

[entlein]

Summary: PEM direct-query gRPC endpoint (entlein/dx#29). Make the metadata-connected vizier-pem serve api.vizierpb.VizierService.ExecuteScript directly over gRPC, JWT-authenticated, so dx queries the node-local PEM with no broker hop. Ports the standalone_pem capability with two upgrades: metadata-connected reuse of the live PEM Carnot + agent metadata (closes #15), and HS256 service-token auth via the cluster jwt-signing-key (no longer insecure). Drain fix (execution_and_timing_info -> QueryData.execution_stats wire-roundtrip + skip payload-less responses) closes the unimplemented type stream error caught on live PEM soak. Fail-soft direct-query startup (try/catch + step 1/6 to 6/6 breadcrumbs) so init failure can never crashloop the data plane. Manager::Init refuses to start with empty PL_JWT_SIGNING_KEY (catches the jwt::SigningError that crashloop'd the stock 0.14.17 PEM).

Relevant Issues: entlein/dx#29 (PEM direct-query)

Type of change: /kind feature

Test Plan: bazel test //src/vizier/services/agent/pem:direct_query_server_test (auth-negative + ValidToken_TrivialQuery_StreamsRows green); bazel test //src/vizier/services/agent/shared/manager/... (5/5 pass, JWT guard wired through Manager::Init); vizier-release CI builds pemdq6 image (commit 50dffb0 / tag release/vizier/v0.14.19-pemdq6); live PG soak

[coderabbitai]

Warning

Review limit reached

@entlein, we couldn't start this review because you've reached your PR review rate limit.

More reviews will be available in 54 minutes and 57 seconds. Learn how PR review limits work.

Your organization has used up its prepaid credits, and credit purchases are no longer available. Enable the review add-on in the billing tab to keep reviews running — you're only billed for reviews past your plan's rate limits ($0.25/file).

⌛ How to resolve this issue?

After more reviews become available, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based credits.

🚦 How do rate limits work?

CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan refill rate.

For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, the refill rate gradually slows as usage increases. The highest same-day bursts are limited more strictly.

Please see our Fair Usage Limits Policy for further information.

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: ASSERTIVE

Plan: Pro Plus

Run ID: f8d45aaa-ed68-49cf-9871-b80c5c26bf56

📥 Commits

Reviewing files that changed from the base of the PR and between e075700 and c357699.

📒 Files selected for processing (18)

• .github/workflows/vizier_release.yaml
• src/api/go/pxapi/opts.go
• src/carnot/BUILD.bazel
• src/carnot/exec/BUILD.bazel
• src/carnot/udf/BUILD.bazel
• src/stirling/source_connectors/socket_tracer/testing/container_images/BUILD.bazel
• src/utils/shared/k8s/apply.go
• src/utils/shared/k8s/delete.go
• src/vizier/services/agent/pem/BUILD.bazel
• src/vizier/services/agent/pem/DIRECT_QUERY_CONTRACT.md
• src/vizier/services/agent/pem/DIRECT_QUERY_SECURITY.md
• src/vizier/services/agent/pem/direct_query_server.cc
• src/vizier/services/agent/pem/direct_query_server.h
• src/vizier/services/agent/pem/direct_query_server_test.cc
• src/vizier/services/agent/pem/pem_main.cc
• src/vizier/services/agent/pem/pem_manager.cc
• src/vizier/services/agent/pem/pem_manager.h
• src/vizier/services/agent/shared/manager/manager.cc

📝 Walkthrough

Walkthrough

Adds a PEM direct-query ExecuteScript gRPC endpoint with HS256 JWT bearer-token verification, server-streaming query execution, comprehensive in-process tests, PEM runtime wiring, Bazel visibility and dependency updates, and minor utility/CI tweaks.

Changes

PEM Direct-Query Feature

Layer / File(s)	Summary
Direct-query contract and interface src/vizier/services/agent/pem/DIRECT_QUERY_CONTRACT.md, src/vizier/services/agent/pem/direct_query_server.h	Adds contract doc specifying gRPC TLS, feature flag gating, JWT authentication, server-streaming with mutation rejection, and dx-agent integration. Header declares AuthenticateRequest helper returning OK only for valid HS256 tokens and DirectQueryServer class with public constructor taking carnot::Carnot*, carnot::EngineState*, LocalGRPCResultSinkServer*, and JWT signing key; ExecuteScript RPC declared for streaming responses.
Server implementation, auth, and streaming src/vizier/services/agent/pem/direct_query_server.cc	Implements case-insensitive Bearer token extraction, RFC 7515 base64url decoding, HS256 HMAC-SHA256 signature verification with constant-time comparison, and claim validation (aud string/array, iss, numeric exp checking, required Scopes with service scope). AuthenticateRequest fail-closed returns UNAUTHENTICATED with VLOG details on failure. ExecuteScript authenticates first, rejects mutations with UNIMPLEMENTED, validates Carnot wiring, compiles PxL, emits schema-only responses, serializes execution with mutex, executes query, and streams drained sink results; includes disabled-build stubs for compile-time kill switch.
In-process contract and execution tests src/vizier/services/agent/pem/direct_query_server_test.cc	Adds MakeBearerToken minting HS256 tokens with claim variants (valid, expired, wrong key, algorithm confusion, malformed aud/iss/exp, missing scope). DirectQueryServerTest fixture runs in-process server with no Carnot; tests assert UNAUTHENTICATED for missing/wrong/tampered tokens and UNIMPLEMENTED for mutations. DirectQueryServerExecTest builds real Carnot/local execution environment with inline schemas, wires gRPC server, and tests end-to-end streaming, concurrency, sequential reuse, broker independence, feature-toggle short-circuit, and robustness (invalid PxL, unknown tables).
PEM runtime wiring and flags src/vizier/services/agent/pem/pem_manager.cc, src/vizier/services/agent/pem/pem_manager.h, src/vizier/services/agent/pem/pem_main.cc, src/vizier/services/agent/shared/manager/manager.cc	Adds gflags (direct_query_enabled, direct_query_port, direct_query_jwt_signing_key) gated by compile-time flag. PEMManager::MaybeStartDirectQueryServer constructs local Carnot/UDF, registers DirectQueryServer, starts gRPC server on 0.0.0.0:<port> via SSL::DefaultGRPCServerCreds; fail-soft on errors. PEMManager::StopDirectQueryServer shuts down server and resets objects. Lifecycle hooks wired to PostRegisterHookImpl and StopImpl. Manager startup validates JWT signing key non-empty.
Bazel build, test target, and visibility src/vizier/services/agent/pem/BUILD.bazel, src/carnot/BUILD.bazel, src/carnot/exec/BUILD.bazel, src/carnot/udf/BUILD.bazel, src/vizier/services/agent/pem/DIRECT_QUERY_SECURITY.md	Adds config_setting(name = "direct_query_disabled", ...) and updates pl_cc_library with conditional defines and expanded deps (Vizier proto, //src/carnot headers/impl, @boringssl//:crypto, gRPC++, rapidjson, sole). Adds pl_cc_test(name = "direct_query_server_test", ...) target. Exports local_grpc_result_server.h from Carnot exec. Expands visibility labels to allow //src/vizier/services/agent/pem:__pkg__ on Carnot modules. Adds detailed security spec documenting signing-key flow, threat model, and test-locked tampering expectations.
K8s utilities, client options, and CI runner updates src/utils/shared/k8s/apply.go, src/utils/shared/k8s/delete.go, src/stirling/source_connectors/socket_tracer/testing/container_images/BUILD.bazel, .github/workflows/vizier_release.yaml, src/api/go/pxapi/opts.go	Reorders k8s import (k8serrors before meta); migrates sets.String → sets.Set[string] usages and set constructors; reorders go_container_libraries Bazel arguments (no functional change); adds WithDirectTLSSkipVerify() client option for direct endpoint TLS cert skipping; updates vizier_release runs-on runners to oracle-vm-16cpu-64gb-x86-64.

Sequence Diagram

sequenceDiagram
  participant Client
  participant DirectQueryServer
  participant JWTVerifier
  participant Carnot
  participant ResultSink
  Client->>DirectQueryServer: ExecuteScriptRequest with Bearer token
  DirectQueryServer->>JWTVerifier: AuthenticateRequest(authorization header)
  JWTVerifier->>JWTVerifier: Extract and base64url decode JWT parts
  JWTVerifier->>JWTVerifier: Verify HS256 HMAC-SHA256 signature
  JWTVerifier->>JWTVerifier: Validate claims (aud, iss, exp, Scopes)
  alt Auth failure
    JWTVerifier-->>DirectQueryServer: UNAUTHENTICATED
    DirectQueryServer-->>Client: UNAUTHENTICATED
  else Auth success
    JWTVerifier-->>DirectQueryServer: OK
    DirectQueryServer->>Carnot: Compile PxL query
    alt Compile failure
      Carnot-->>DirectQueryServer: error
      DirectQueryServer-->>Client: INVALID_ARGUMENT
    else Compile success
      Carnot-->>DirectQueryServer: compiled plan with sinks
      DirectQueryServer->>ResultSink: ResetQueryResults()
      DirectQueryServer->>Carnot: ExecuteQuery()
      loop Drain results
        ResultSink-->>DirectQueryServer: ExecuteScriptResponse (schema/rows/stats)
        DirectQueryServer-->>Client: stream response
      end
      DirectQueryServer-->>Client: OK (end stream)
    end
  end

Loading

🎯 3 (Moderate) | ⏱️ ~25 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 71.60% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title accurately reflects the main change: adding a direct-query gRPC endpoint to PEM with JWT authentication and comprehensive test coverage.
Description check	✅ Passed	The description is clearly related to the changeset, explaining the PEM direct-query gRPC endpoint implementation with JWT authentication, configuration, and testing.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch entlein/pem-direct-query

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@src/vizier/services/agent/pem/direct_query_server_test.cc`:
- Around line 58-67: The test ValidToken_Mutation_Unimplemented is unreachable
because MakeBearerToken returns placeholder non-JWT strings and
AuthenticateRequest always returns UNAUTHENTICATED; change the test surface so a
"valid" token is actually recognized: either make MakeBearerToken produce a
token format that AuthenticateRequest will accept as valid (e.g., generate a
simple signed JWT or a recognized test token when TokenKind::kValid) or adjust
AuthenticateRequest (in direct_query_server.cc's AuthenticateRequest) to accept
the test placeholder for TokenKind::kValid; update MakeBearerToken's TokenKind
cases (kValid, kWrongKey, kExpired) to return distinct tokens that map to
AuthenticateRequest's logic so the ValidToken_Mutation_Unimplemented path can
observe UNIMPLEMENTED as intended.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: ASSERTIVE

Plan: Pro Plus

Run ID: 649f9e15-2555-4317-9973-b200df724e3d

📥 Commits

Reviewing files that changed from the base of the PR and between 9b2721f and 7c717c5.

📒 Files selected for processing (5)

• src/vizier/services/agent/pem/BUILD.bazel
• src/vizier/services/agent/pem/DIRECT_QUERY_CONTRACT.md
• src/vizier/services/agent/pem/direct_query_server.cc
• src/vizier/services/agent/pem/direct_query_server.h
• src/vizier/services/agent/pem/direct_query_server_test.cc

[coderabbitai]

Actionable comments posted: 1

♻️ Duplicate comments (1)

src/vizier/services/agent/pem/direct_query_server.cc (1)

175-218: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Still missing the cluster-service claim checks.

This verifier still accepts any HS256 token signed with the shared key and aud=="vizier". That is broader than the stated cluster-service JWT contract, because iss, sub, Scopes, and ServiceID are never enforced here.

Suggested tightening

   if (!aud_ok) {
     return ::grpc::Status(::grpc::StatusCode::UNAUTHENTICATED,
                           "direct-query: wrong audience (expected vizier)");
   }
+  if (!payload.HasMember("iss") || !payload["iss"].IsString() ||
+      std::strcmp(payload["iss"].GetString(), "PL") != 0 ||
+      !payload.HasMember("sub") || !payload["sub"].IsString() ||
+      std::strcmp(payload["sub"].GetString(), "service") != 0 ||
+      !payload.HasMember("Scopes") || !payload["Scopes"].IsString() ||
+      std::strcmp(payload["Scopes"].GetString(), "service") != 0 ||
+      !payload.HasMember("ServiceID") || !payload["ServiceID"].IsString() ||
+      payload["ServiceID"].GetStringLength() == 0) {
+    return ::grpc::Status(::grpc::StatusCode::UNAUTHENTICATED,
+                          "direct-query: not a cluster service JWT");
+  }
   if (!payload.HasMember("exp")) {
     return ::grpc::Status(::grpc::StatusCode::UNAUTHENTICATED, "direct-query: missing exp claim");
   }

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/vizier/services/agent/pem/direct_query_server.cc` around lines 175 - 218,
The verifier currently only checks aud and exp but must enforce the
cluster-service JWT contract: after the aud check and before exp validation, add
explicit checks on payload for "iss" (must be string == kExpectedIssuer), "sub"
(must be string == kExpectedSubject), "Scopes" (must contain the cluster-service
scope, e.g., check string or array includes kClusterServiceScope), and
"ServiceID" (must exist and match the expected service/cluster id, e.g.,
kExpectedServiceID or the runtime cluster id). For each missing or mismatched
claim return ::grpc::Status(::grpc::StatusCode::UNAUTHENTICATED, "<contextual
message>") similar to the existing messages; use the existing payload variable
and keep the new constants kExpectedIssuer, kExpectedSubject,
kClusterServiceScope, kExpectedServiceID (or existing equivalents) to locate
where to enforce these checks.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@src/vizier/services/agent/pem/direct_query_server_test.cc`:
- Around line 196-236: The test uses an empty http_events table so it never
exercises the drainSinkAndStream() column-copy path; seed at least one row into
the table (e.g., in MakeHTTPEventsTable() or in
TEST_F(DirectQueryServerExecTest, ValidToken_TrivialQuery_StreamsRows) before
calling stub_->ExecuteScript) using the Table API for "http_events" (populate
fields time_, upid, remote_addr, remote_port, trace_role), then when streaming
responses from stub_->ExecuteScript assert that you observe a streamed batch
with resp.has_table_id() == true (or matches the expected table id) and
resp.num_rows() > 0 to validate the non-empty row-batch path is exercised.

---

Duplicate comments:
In `@src/vizier/services/agent/pem/direct_query_server.cc`:
- Around line 175-218: The verifier currently only checks aud and exp but must
enforce the cluster-service JWT contract: after the aud check and before exp
validation, add explicit checks on payload for "iss" (must be string ==
kExpectedIssuer), "sub" (must be string == kExpectedSubject), "Scopes" (must
contain the cluster-service scope, e.g., check string or array includes
kClusterServiceScope), and "ServiceID" (must exist and match the expected
service/cluster id, e.g., kExpectedServiceID or the runtime cluster id). For
each missing or mismatched claim return
::grpc::Status(::grpc::StatusCode::UNAUTHENTICATED, "<contextual message>")
similar to the existing messages; use the existing payload variable and keep the
new constants kExpectedIssuer, kExpectedSubject, kClusterServiceScope,
kExpectedServiceID (or existing equivalents) to locate where to enforce these
checks.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: ASSERTIVE

Plan: Pro Plus

Run ID: 325769b9-0ed7-4ca0-8704-52d0312dd1e0

📥 Commits

Reviewing files that changed from the base of the PR and between 00e198c and b409464.

📒 Files selected for processing (10)

• .github/workflows/vizier_release.yaml
• src/carnot/BUILD.bazel
• src/carnot/exec/BUILD.bazel
• src/carnot/udf/BUILD.bazel
• src/vizier/services/agent/pem/BUILD.bazel
• src/vizier/services/agent/pem/direct_query_server.cc
• src/vizier/services/agent/pem/direct_query_server_test.cc
• src/vizier/services/agent/pem/pem_main.cc
• src/vizier/services/agent/pem/pem_manager.cc
• src/vizier/services/agent/pem/pem_manager.h

[coderabbitai]

Actionable comments posted: 1

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

src/vizier/services/agent/pem/direct_query_server.cc (1)

411-422: ⚠️ Potential issue | 🔴 Critical | ⚡ Quick win

Serialize sink reset/execute/drain to prevent cross-request result corruption.

result_server_ is shared state, and this method resets plus drains global accumulated chunks. Concurrent ExecuteScript calls can clobber each other and stream mixed results.

Suggested minimal guard (serialize direct-query execution path)

+#include "absl/synchronization/mutex.h"
...
 namespace {
+absl::Mutex g_direct_query_exec_mu;
 }  // namespace
...
 ::grpc::Status DirectQueryServer::ExecuteScript(
     ::grpc::ServerContext* context, const ::px::api::vizierpb::ExecuteScriptRequest* request,
     ::grpc::ServerWriter<::px::api::vizierpb::ExecuteScriptResponse>* writer) {
+  absl::MutexLock lk(&g_direct_query_exec_mu);
   if (auto s = AuthenticateRequest(context, jwt_signing_key_); !s.ok()) {
     return s;
   }
...
   result_server_->ResetQueryResults();
   auto exec_s = carnot_->ExecuteQuery(request->query_str(), query_id, ::px::CurrentTimeNS());

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/vizier/services/agent/pem/direct_query_server.cc` around lines 411 - 422,
This RPC resets and reads the shared result sink (result_server_) around
carnot_->ExecuteQuery and drainSinkAndStream, so concurrent
ExecuteScript/ExecuteQuery calls can interleave and corrupt results; serialize
the sequence by introducing a mutex (e.g., a class member direct_query_mu_) and
acquire a std::lock_guard (or std::unique_lock) at the start of the method that
surrounds result_server_->ResetQueryResults(), the call to
carnot_->ExecuteQuery(...), and drainSinkAndStream(result_server_, query_id_str,
writer) so the reset/execute/drain happens atomically for a single request.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@src/vizier/services/agent/pem/pem_manager.cc`:
- Around line 39-42: The code currently uses FLAGS_direct_query_jwt_signing_key
separately from FLAGS_jwt_signing_key causing split-brain; change the logic in
pem_manager.cc to treat FLAGS_direct_query_jwt_signing_key as optional and fall
back to FLAGS_jwt_signing_key when empty (i.e., wherever
FLAGS_direct_query_jwt_signing_key is read—references around the DEFINE_string
and the usages near the blocks you noted at lines ~121-125 and ~161-163—use a
single effective key variable like effective_direct_query_key =
FLAGS_direct_query_jwt_signing_key.empty() ? FLAGS_jwt_signing_key :
FLAGS_direct_query_jwt_signing_key and use that variable for direct-query
auth/minting).

---

Outside diff comments:
In `@src/vizier/services/agent/pem/direct_query_server.cc`:
- Around line 411-422: This RPC resets and reads the shared result sink
(result_server_) around carnot_->ExecuteQuery and drainSinkAndStream, so
concurrent ExecuteScript/ExecuteQuery calls can interleave and corrupt results;
serialize the sequence by introducing a mutex (e.g., a class member
direct_query_mu_) and acquire a std::lock_guard (or std::unique_lock) at the
start of the method that surrounds result_server_->ResetQueryResults(), the call
to carnot_->ExecuteQuery(...), and drainSinkAndStream(result_server_,
query_id_str, writer) so the reset/execute/drain happens atomically for a single
request.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: ASSERTIVE

Plan: Pro Plus

Run ID: 271756a5-b762-4b8d-9811-f5b171d98116

📥 Commits

Reviewing files that changed from the base of the PR and between b409464 and 50dffb0.

📒 Files selected for processing (4)

• src/vizier/services/agent/pem/direct_query_server.cc
• src/vizier/services/agent/pem/pem_main.cc
• src/vizier/services/agent/pem/pem_manager.cc
• src/vizier/services/agent/shared/manager/manager.cc

[coderabbitai]

Actionable comments posted: 1

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

src/vizier/services/agent/pem/direct_query_server.cc (1)

301-302: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Stop work when stream writes fail.

In src/vizier/services/agent/pem/direct_query_server.cc, the ServerWriter::Write(...) return values are ignored at writer->Write(schema_resp); (around lines 301) and writer->Write(resp); (around line 373). If the client disconnects, the server can keep doing schema/result processing instead of aborting early.

Handle Write(...) failures by propagating them up and returning a CANCELLED/early-abort status to stop further work/draining.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/vizier/services/agent/pem/direct_query_server.cc` around lines 301 - 302,
The server currently ignores the boolean return from writer->Write(...) calls
(notably the calls with schema_resp and resp in direct_query_server.cc), so if
the client disconnects the server continues processing; modify the enclosing
methods (the RPC handler or helper functions around the
writer->Write(schema_resp) and writer->Write(resp) sites) to check the
Write(...) return value and, on failure, immediately stop further work and
return a gRPC CANCELLED status (or propagate a failure status) up to the caller
so processing/draining aborts; ensure any callers of those helpers propagate
that Status instead of ignoring it.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@src/vizier/services/agent/pem/direct_query_server.cc`:
- Around line 329-374: This code leaks shared mutable sink state across
concurrent ExecuteScript calls because ResetQueryResults(), ExecuteQuery(), and
draining raw_query_results() access the same sink; protect the critical section
by serializing per-request access (e.g., use a mutex or request-scoped lock)
around the sequence that calls ResetQueryResults(), ExecuteQuery(), and the loop
over result_server->raw_query_results() so chunks cannot be cleared or
interleaved by another request; apply the same protection to the analogous block
referenced at the other location (the 414-421 section) and ensure the lock is
held from before ResetQueryResults() until after the writer->Write(resp) loop
completes.

---

Outside diff comments:
In `@src/vizier/services/agent/pem/direct_query_server.cc`:
- Around line 301-302: The server currently ignores the boolean return from
writer->Write(...) calls (notably the calls with schema_resp and resp in
direct_query_server.cc), so if the client disconnects the server continues
processing; modify the enclosing methods (the RPC handler or helper functions
around the writer->Write(schema_resp) and writer->Write(resp) sites) to check
the Write(...) return value and, on failure, immediately stop further work and
return a gRPC CANCELLED status (or propagate a failure status) up to the caller
so processing/draining aborts; ensure any callers of those helpers propagate
that Status instead of ignoring it.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: ASSERTIVE

Plan: Pro Plus

Run ID: cfc25291-1a5a-4a5a-8032-7380f3e88e1b

📥 Commits

Reviewing files that changed from the base of the PR and between b409464 and 50dffb0.

📒 Files selected for processing (4)

• src/vizier/services/agent/pem/direct_query_server.cc
• src/vizier/services/agent/pem/pem_main.cc
• src/vizier/services/agent/pem/pem_manager.cc
• src/vizier/services/agent/shared/manager/manager.cc

[entlein]

Review for claude-build-agent:

• exlicitely state the new contracts
• the fallback scenarios need explicit tests
• the new security exposure needs to be explained in a Readme -> how do clients authenticate (and what processes are discouraged and WHY)
• test the various failure modes in case the authentication doesnt suceed
• create a apples-to-apples benchmark test for pem (upstream) vs dual-usage-pem (this PR), clearly profile the root causes of any discrepancies and in case of discovered tech-debt, post the numbers and triage to an issue.
• connecting locally to pem is an additional attack surface, and it must be possible to fully disable it. create evidence for the feature-toggle being 100% effective in case the feature is not desired.
• there should be a compiler flag that fully disables the feature in case customers do not want the feature available in the binary

Kind regard, your human user (I am not the pixie-agent)
And: this is pre-requisite to merging.

[entlein]

please address the insecure credentials gap: this is a blocker. implement now

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@src/vizier/services/agent/pem/DIRECT_QUERY_SECURITY.md`:
- Around line 1-332: The doc's cross-reference to the DEFINE location is out of
date: update the mention of pem_manager.cc:39 to the current line
(pem_manager.cc:47) where FLAGS_direct_query_jwt_signing_key's DEFINE_string now
lives; search for the text "pem_manager.cc:39" in DIRECT_QUERY_SECURITY.md and
replace it with "pem_manager.cc:47" and run a quick grep to ensure no other
stale numeric references to pem_manager.cc remain, keeping the reference text
around the symbol FLAGS_direct_query_jwt_signing_key unchanged.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: ASSERTIVE

Plan: Pro Plus

Run ID: 4486db87-17b2-44c9-9f62-57db22524ab2

📥 Commits

Reviewing files that changed from the base of the PR and between 66d92c5 and 847409f.

📒 Files selected for processing (6)

• src/vizier/services/agent/pem/BUILD.bazel
• src/vizier/services/agent/pem/DIRECT_QUERY_SECURITY.md
• src/vizier/services/agent/pem/direct_query_server.cc
• src/vizier/services/agent/pem/direct_query_server.h
• src/vizier/services/agent/pem/direct_query_server_test.cc
• src/vizier/services/agent/pem/pem_manager.cc

[entlein]

@CodeRabbit : rereview the entire PR now, it seems ready to merge

[coderabbitai]

Tip

For best results, initiate chat on the files or code changes.

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
echo "=== direct_query_server.h ==="
cat src/vizier/services/agent/pem/direct_query_server.h

echo ""
echo "=== direct_query_server.cc ==="
cat src/vizier/services/agent/pem/direct_query_server.cc

echo ""
echo "=== direct_query_server_test.cc (first 300 lines) ==="
head -300 src/vizier/services/agent/pem/direct_query_server_test.cc

Length of output: 41934

---

🏁 Script executed:

#!/bin/bash
echo "=== direct_query_server_test.cc (lines 300-600) ==="
sed -n '300,600p' src/vizier/services/agent/pem/direct_query_server_test.cc

echo ""
echo "=== direct_query_server_test.cc (lines 600-895) ==="
sed -n '600,895p' src/vizier/services/agent/pem/direct_query_server_test.cc

Length of output: 30332

---

🏁 Script executed:

#!/bin/bash
echo "=== pem_manager.cc (new additions only - last 200 lines) ==="
tail -250 src/vizier/services/agent/pem/pem_manager.cc

echo ""
echo "=== pem_manager.h ==="
cat src/vizier/services/agent/pem/pem_manager.h

echo ""
echo "=== DIRECT_QUERY_SECURITY.md ==="
cat src/vizier/services/agent/pem/DIRECT_QUERY_SECURITY.md

echo ""
echo "=== DIRECT_QUERY_CONTRACT.md ==="
cat src/vizier/services/agent/pem/DIRECT_QUERY_CONTRACT.md

Length of output: 44645

---

🏁 Script executed:

#!/bin/bash
echo "=== BUILD.bazel (pem) ==="
cat src/vizier/services/agent/pem/BUILD.bazel

echo ""
echo "=== manager.cc changes ==="
cat src/vizier/services/agent/shared/manager/manager.cc | grep -A 20 "jwt_signing_key"

echo ""
echo "=== opts.go ==="
cat src/api/go/pxapi/opts.go

echo ""
echo "=== delete.go ==="
cat src/utils/shared/k8s/delete.go

Length of output: 25624

---

---

@entlein — comprehensive re-review completed. Here is the full assessment:

---

Overall: Functionally merge-ready ✅

All seven original review requirements and both blockers (TLS + auth bug) are resolved. The implementation is solid. I found three minor issues below — none block merge, but one has a misleading comment worth correcting.

---

✅ What is now correct

Area	Status
Auth bug (sub=="service" → Scopes check)	Fixed in verifyHs256Jwt; unit tests mint realistic tokens (sub="dx", Scopes="service")
TLS (InsecureServerCredentials → SSL::DefaultGRPCServerCreds())	Fixed in MaybeStartDirectQueryServer (step 6/6)
Security documentation (DIRECT_QUERY_SECURITY.md)	Comprehensive — key flow, threat model, transport, discouraged practices, failure modes, rotation limitations
Feature toggles (runtime + compile-time)	--direct_query_enabled=false skips port binding; PX_PEM_DIRECT_QUERY=disabled strips all feature bytes
Compile-time toggle tests	CompiledOut_ValidToken_StillUnauthenticated + ToggleContract_DocumentBothLevels cover both modes
JWT tampering coverage	18+ named test cases covering alg-none, sig-flip, payload-flip, truncation, concatenation, alg-confusion
Fail-soft startup	6-step breadcrumbs, all failure paths return Status::OK() to avoid crashing the data plane
Scope check (service tokens)	Scopes claim split on , and checked for "service"; not subject claim
Mutex serialization	exec_mu_ per-instance, guards reset/execute/drain; ConcurrentQueries_AllSucceed tests N=8 parallel callers
Wire-format roundtrip drain	drainSinkAndStream skips payload-less chunks; execution_and_timing_info → execution_stats mapped

---

🟡 Minor issues (non-blocking)

1. Misleading comment in constantTimeEquals

direct_query_server.cc, the comment on constantTimeEquals says:

"Mismatched-length inputs trivially differ but we still walk the shorter to keep timing predictable across malformed lengths."

But the code returns immediately:

bool constantTimeEquals(absl::string_view a, absl::string_view b) {
  if (a.size() != b.size()) return false;  // ← returns immediately, doesn't walk
  ...
}

In practice this is not a security concern here since computed_mac (HMAC-SHA256) is always exactly 32 bytes, so the length branch is only hit when the decoded signature has a different length — which already means the token is malformed. Still, the comment should match the code. Either remove the "we still walk" sentence or add the fixed-length early-exit path before the size check to make the comment truthful.

2. exp claim doesn't handle int32

direct_query_server.cc:228–238 handles IsInt64, IsUint64, IsDouble but not IsInt (RapidJSON's 32-bit integer type). JWT libraries can technically serialize exp as a 32-bit integer for dates before year 2038. This will be hit only on non-standard JWT minters but is a one-liner fix:

} else if (payload["exp"].IsInt()) {
  exp_secs = static_cast<int64_t>(payload["exp"].GetInt());
}

3. DIRECT_QUERY_CONTRACT.md env-var label for the direct-query signing key

Line 20 of DIRECT_QUERY_CONTRACT.md maps --direct_query_jwt_signing_key to env PL_JWT_SIGNING_KEY, which is the same env var name as the main signing key in manager.cc. The actual flag definition in pem_manager.cc almost certainly reads a different env var (e.g. PL_DIRECT_QUERY_JWT_SIGNING_KEY or falls back to PL_JWT_SIGNING_KEY). If the env var really is shared, the table is accurate but could confuse operators who try to set one independently. Worth a clarifying note in the doc.

---

Remaining skipped tests (acceptable)

• FailSoft_BrokerFailureToleratedByDirectQuery — SKIP'd with a clear explanation and a tracked follow-up (hoist MaybeStartDirectQueryServer or broker-optional Manager mode). Acceptable for this PR's scope.
• Benchmark_PemDirect_Vs_BrokerPath_RedPlaceholder — SKIP'd with pemdq5 soak data cited (43.5s/q vs 27s/q, dominant factor = second Carnot). The placeholder names the issue greppably.
• PerPodFilter_MetadataConnected — SKIP'd with correct rationale (needs metadata + multi-pod fixture).

---

Verdict

The implementation is complete. The critical auth bug and TLS blocker that were merge-blockers are both resolved and live-verified on pemdq10. The three items above are cosmetic/defensive and can be addressed in a follow-up. This PR is ready to merge.