Skip to content

feat: implement ptq rpc and index schemas #3138

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

Draft
wilsonrivera wants to merge 4 commits into from
Draft

feat: implement ptq rpc and index schemas #3138

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
b9c52bf
feat: implement ptq rpc and index schemas
wilsonrivera Aug 5, 2026
8f2fc5a
chore: remove `console.log`
wilsonrivera Aug 5, 2026
0a70606
chore: fix router test
wilsonrivera Aug 5, 2026
c9d6d43
Merge branch 'main' into wilson/cosmo-335-cosmo-cloud-prompt-to-query…
wilsonrivera Aug 5, 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
1,104 changes: 681 additions & 423 deletions connect-go/gen/proto/wg/cosmo/node/v1/node.pb.go
View file
Open in desktop

Large diffs are not rendered by default.

31 changes: 30 additions & 1 deletion connect-go/gen/proto/wg/cosmo/node/v1/nodev1connect/node.connect.go
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

5 changes: 5 additions & 0 deletions connect/src/wg/cosmo/node/v1/node-NodeService_connectquery.ts
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

255 changes: 185 additions & 70 deletions connect/src/wg/cosmo/node/v1/node_pb.ts
View file
Open in desktop

Large diffs are not rendered by default.

3 changes: 3 additions & 0 deletions controlplane/.env.example
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -85,3 +85,6 @@ SENTRY_PROFILE_SESSION_SAMPLE_RATE="1.0"
SENTRY_PROFILE_LIFECYCLE="manual"
SENTRY_EVENT_LOOP_BLOCK_THRESHOLD_MS="100"
SENTRY_ENABLE_LOGS="false"

# Prompt To Query
PROMPT_TO_QUERY_SERVICE_ENDPOINT=
21 changes: 20 additions & 1 deletion controlplane/src/core/bufservices/NodeService.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,12 +1,13 @@
import { ServiceImpl } from '@connectrpc/connect';
import { EnumStatusCode } from '@wundergraph/cosmo-connect/dist/common/common_pb';
import { NodeService, RegistrationInfo, SelfRegisterResponse } from '@wundergraph/cosmo-connect/dist/node/v1/node_pb';
import { NodeService, RegistrationInfo } from '@wundergraph/cosmo-connect/dist/node/v1/node_pb';
import { lru } from 'tiny-lru';
import { FederatedGraphRepository } from '../repositories/FederatedGraphRepository.js';
import { OrganizationRepository } from '../repositories/OrganizationRepository.js';
import type { PlainMessage } from '../../types/index.js';
import type { RouterOptions } from '../routes.js';
import { enrichLogger, getLogger, handleError } from '../util.js';
import { PromptToQueryService } from '../services/PromptToQueryService.js';

export default function (opts: RouterOptions): Partial<ServiceImpl<typeof NodeService>> {
const registrationInfoCache = lru<PlainMessage<RegistrationInfo>>(1000, 300_000);
Expand Down Expand Up @@ -66,5 +67,23 @@ export default function (opts: RouterOptions): Partial<ServiceImpl<typeof NodeSe
};
});
},

generateQuery(req, ctx) {
let logger = getLogger(ctx, opts.logger);
return handleError(ctx, logger, async () => {
const authContext = await opts.authenticator.authenticateRouter(ctx.requestHeader);
logger = enrichLogger(ctx, logger, authContext);

const ptqService = new PromptToQueryService(
opts.db,
logger,
opts.promptToQueryServiceAddress,
authContext.organizationId,
opts.billingDefaultPlanId,
);

return ptqService.generateQuery(req.schemaHash, req.prompt);

@coderabbitai coderabbitai Bot Aug 5, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔒 Security & Privacy | 🟠 Major | 🏗️ Heavy lift

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
set -euo pipefail

ast-grep outline controlplane/src/core/bufservices/NodeService.ts --items all

rg -n -C 5 'schemaHash|schemaSha|indexId|EnsureIndex|indexSchema|GenerateQuery' controlplane connect proto
rg -n -C 5 'authenticateRouter|federatedGraphId|organizationId' controlplane/src

Repository: wundergraph/cosmo

Length of output: 50374

Authorization Bypass (CWE-639): Authorization Bypass Through User-Controlled Key (IDOR)

Reachability: External · Exploitability: Moderate

Reachability path 
● Entry
  connect/src/wg/cosmo/node/v1/node-NodeService_connectquery.ts:15
  generateQuery
│
▼
● Sink
  controlplane/src/core/bufservices/NodeService.ts

Bind schemaHash to the authenticated router graph.

generateQuery scopes only the feature check to authContext.organizationId and forwards the caller-controlled hash as indexId. Resolve and authorize the hash for authContext.federatedGraphId and authContext.organizationId before calling PromptToQueryService.

🤖 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 `@controlplane/src/core/bufservices/NodeService.ts` at line 85, Update the
request flow in NodeService around ptqService.generateQuery to resolve the
requested schemaHash against authContext.federatedGraphId and
authContext.organizationId, authorize that association, and pass only the
validated graph-bound hash to PromptToQueryService instead of forwarding the
caller-controlled value directly.

});
},
};
}
2 changes: 2 additions & 0 deletions controlplane/src/core/bufservices/contract/createContract.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -182,6 +182,8 @@ export function createContract(
opts.chClient,
opts.webhookProxyUrl,
req.disableResolvabilityValidation,
opts.promptToQueryServiceAddress,
opts.billingDefaultPlanId,
);

const { deploymentErrors, compositionErrors, compositionWarnings } =
Expand Down
2 changes: 2 additions & 0 deletions controlplane/src/core/bufservices/contract/updateContract.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -109,6 +109,8 @@ export function updateContract(
opts.chClient,
opts.webhookProxyUrl,
req.disableResolvabilityValidation,
opts.promptToQueryServiceAddress,
opts.billingDefaultPlanId,
);

// Update the contract details
Expand Down
2 changes: 2 additions & 0 deletions controlplane/src/core/bufservices/feature-flag/createFeatureFlag.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -198,6 +198,8 @@ export function createFeatureFlag(
opts.chClient,
opts.webhookProxyUrl,
req.disableResolvabilityValidation,
opts.promptToQueryServiceAddress,
opts.billingDefaultPlanId,
);

return compositionService.composeAndDeployFeatureFlag({
Expand Down
3 changes: 2 additions & 1 deletion controlplane/src/core/bufservices/feature-flag/deleteFeatureFlag.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -68,7 +68,6 @@ export function deleteFeatureFlag(

const { deploymentErrors, compositionErrors, compositionWarnings } = await opts.db.transaction(async (tx) => {
const auditLogRepo = new AuditLogRepository(tx);
const featureFlagRepo = new FeatureFlagRepository(logger, tx, authContext.organizationId);
const compositionService = new CompositionService(
tx,
authContext.organizationId,
Expand All @@ -78,6 +77,8 @@ export function deleteFeatureFlag(
opts.chClient,
opts.webhookProxyUrl,
req.disableResolvabilityValidation,
opts.promptToQueryServiceAddress,
opts.billingDefaultPlanId,
);

const result = await compositionService.deleteFeatureFlag({
Expand Down
2 changes: 2 additions & 0 deletions controlplane/src/core/bufservices/feature-flag/enableFeatureFlag.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -96,6 +96,8 @@ export function enableFeatureFlag(
opts.chClient,
opts.webhookProxyUrl,
req.disableResolvabilityValidation,
opts.promptToQueryServiceAddress,
opts.billingDefaultPlanId,
);

return compositionService.composeAndDeployFeatureFlag({
Expand Down
2 changes: 2 additions & 0 deletions controlplane/src/core/bufservices/feature-flag/recomposeFeatureFlag.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -100,6 +100,8 @@ export function recomposeFeatureFlag(
opts.chClient,
opts.webhookProxyUrl,
req.disableResolvabilityValidation,
opts.promptToQueryServiceAddress,
opts.billingDefaultPlanId,
);

return compositionService.composeAndDeployFeatureFlag({
Expand Down
2 changes: 2 additions & 0 deletions controlplane/src/core/bufservices/feature-flag/updateFeatureFlag.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -203,6 +203,8 @@ export function updateFeatureFlag(
opts.chClient,
opts.webhookProxyUrl,
req.disableResolvabilityValidation,
opts.promptToQueryServiceAddress,
opts.billingDefaultPlanId,
);

const compositionResult = await compositionService.composeAndDeployFeatureFlag({
Expand Down
2 changes: 2 additions & 0 deletions controlplane/src/core/bufservices/federated-graph/createFederatedGraph.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -217,6 +217,8 @@ export function createFederatedGraph(
opts.chClient,
opts.webhookProxyUrl,
req.disableResolvabilityValidation,
opts.promptToQueryServiceAddress,
opts.billingDefaultPlanId,
);

return compositionService.composeAndDeployFederatedGraph({ actorId: authContext.userId, federatedGraph });
Expand Down
2 changes: 2 additions & 0 deletions controlplane/src/core/bufservices/federated-graph/migrateFromApollo.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -157,6 +157,8 @@ export function migrateFromApollo(
opts.chClient,
opts.webhookProxyUrl,
true,
opts.promptToQueryServiceAddress,
opts.billingDefaultPlanId,
);

await compositionService.composeAndDeployFederatedGraph({ actorId: authContext.userId, federatedGraph });
Expand Down
2 changes: 2 additions & 0 deletions controlplane/src/core/bufservices/federated-graph/moveFederatedGraph.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -150,6 +150,8 @@ export function moveFederatedGraph(
opts.chClient,
opts.webhookProxyUrl,
req.disableResolvabilityValidation,
opts.promptToQueryServiceAddress,
opts.billingDefaultPlanId,
);

// Only call `composeAndDeployFederatedGraph` after all contracts have been moved
Expand Down
2 changes: 2 additions & 0 deletions controlplane/src/core/bufservices/federated-graph/updateFederatedGraph.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -115,6 +115,8 @@ export function updateFederatedGraph(
opts.chClient,
opts.webhookProxyUrl,
req.disableResolvabilityValidation,
opts.promptToQueryServiceAddress,
opts.billingDefaultPlanId,
);

const fedGraphRepo = new FederatedGraphRepository(logger, tx, authContext.organizationId);
Expand Down
2 changes: 2 additions & 0 deletions controlplane/src/core/bufservices/graph/recomposeGraph.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,8 @@ export function recomposeGraph(
opts.chClient,
opts.webhookProxyUrl,
req.disableResolvabilityValidation,
opts.promptToQueryServiceAddress,
opts.billingDefaultPlanId,
);

return compositionService.composeAndDeployFederatedGraph({ actorId: authContext.userId, federatedGraph: graph });
Expand Down
2 changes: 2 additions & 0 deletions controlplane/src/core/bufservices/graph/setGraphRouterCompatibilityVersion.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -141,6 +141,8 @@ export function setGraphRouterCompatibilityVersion(
opts.chClient,
opts.webhookProxyUrl,
req.disableResolvabilityValidation,
opts.promptToQueryServiceAddress,
opts.billingDefaultPlanId,
);

return await compositionService.composeAndDeployFederatedGraph({ actorId: authContext.userId, federatedGraph });
Expand Down
2 changes: 2 additions & 0 deletions controlplane/src/core/bufservices/monograph/publishMonograph.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -152,6 +152,8 @@ export function publishMonograph(
opts.chClient,
opts.webhookProxyUrl,
false,
opts.promptToQueryServiceAddress,
opts.billingDefaultPlanId,
);

return subgraphRepo.update(
Expand Down
2 changes: 2 additions & 0 deletions controlplane/src/core/bufservices/monograph/updateMonograph.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -145,6 +145,8 @@ export function updateMonograph(
opts.chClient!,
opts.webhookProxyUrl,
false,
opts.promptToQueryServiceAddress,
opts.billingDefaultPlanId,
);

// Update the federated graph
Expand Down
2 changes: 2 additions & 0 deletions controlplane/src/core/bufservices/subgraph/deleteFederatedSubgraph.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -198,6 +198,8 @@ export function deleteFederatedSubgraph(
opts.chClient,
opts.webhookProxyUrl,
req.disableResolvabilityValidation,
opts.promptToQueryServiceAddress,
opts.billingDefaultPlanId,
);

if (subgraph.isFeatureSubgraph) {
Expand Down
2 changes: 2 additions & 0 deletions controlplane/src/core/bufservices/subgraph/moveSubgraph.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -112,6 +112,8 @@ export function moveSubgraph(
opts.chClient,
opts.webhookProxyUrl,
req.disableResolvabilityValidation,
opts.promptToQueryServiceAddress,
opts.billingDefaultPlanId,
);

const { deploymentErrors, compositionErrors, compositionWarnings, updatedFederatedGraphs } =
Expand Down
2 changes: 2 additions & 0 deletions controlplane/src/core/bufservices/subgraph/publishFederatedSubgraph.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -574,6 +574,8 @@ export function publishFederatedSubgraph(
opts.chClient,
opts.webhookProxyUrl,
req.disableResolvabilityValidation,
opts.promptToQueryServiceAddress,
opts.billingDefaultPlanId,
);

return subgraphRepo.update(
Expand Down
2 changes: 2 additions & 0 deletions controlplane/src/core/bufservices/subgraph/publishFederatedSubgraphs.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -433,6 +433,8 @@ async function runBatchPublish({
opts.chClient,
opts.webhookProxyUrl,
disableResolvabilityValidation,
opts.promptToQueryServiceAddress,
opts.billingDefaultPlanId,
);

const { compositionErrors, compositionWarnings, deploymentErrors } =
Expand Down
2 changes: 2 additions & 0 deletions controlplane/src/core/bufservices/subgraph/updateSubgraph.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -199,6 +199,8 @@ export function updateSubgraph(
opts.chClient,
opts.webhookProxyUrl,
req.disableResolvabilityValidation,
opts.promptToQueryServiceAddress,
opts.billingDefaultPlanId,
);

return subgraphRepo.update(
Expand Down
4 changes: 4 additions & 0 deletions controlplane/src/core/build-server.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -165,6 +165,9 @@ export interface BuildConfig {
key?: string; // e.g. string or '/path/to/my/client-key.pem'
};
};
promptToQuery?: {
address: string | undefined;
};
}

export interface MetricsOptions {
Expand Down Expand Up @@ -624,6 +627,7 @@ export default async function build(opts: BuildConfig) {
webhookProxyUrl: opts.webhook?.proxyUrl,
cdnBaseUrl: opts.cdnBaseUrl,
lockAdapter: fastify.lockAdapter,
promptToQueryServiceAddress: opts.promptToQuery?.address,
}),
contextValues(req) {
const values = createContextValues().set<FastifyBaseLogger>(
Expand Down
4 changes: 4 additions & 0 deletions controlplane/src/core/composition/composer.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@ import { NamespaceRepository } from '../repositories/NamespaceRepository.js';
import { InspectorSchemaChange } from '../services/SchemaUsageTrafficInspector.js';
import { SchemaCheckChangeAction } from '../../db/models.js';
import { traced } from '../tracing.js';
import { PromptToQueryService } from '../services/PromptToQueryService.js';
import {
composeGraphsInWorker,
DeserializedComposedGraph,
Expand Down Expand Up @@ -451,6 +452,7 @@ export class Composer {
routerExecutionConfig,
featureFlagId,
splitConfigEnabled,
promptToQueryService,
}: {
composedGraph: ComposedFederatedGraph;
composedById: string;
Expand All @@ -459,6 +461,7 @@ export class Composer {
routerExecutionConfig?: RouterConfig;
featureFlagId: string;
splitConfigEnabled: boolean;
promptToQueryService: PromptToQueryService;
}): Promise<CompositionDeployResult> {
// For a feature-flag composition the baseline is the previous composition of the same feature flag (not the base
// graph's latest valid version, which would produce a meaningless base-vs-feature-flag diff). Computed before
Expand Down Expand Up @@ -514,6 +517,7 @@ export class Composer {
);
}

await promptToQueryService.indexSchema(composedGraph.composedSchema);
if (schemaChanges.kind !== 'failure' && schemaChanges.changes.length > 0) {
await this.federatedGraphRepo.createFederatedGraphChangelog({
schemaVersionID: updatedFederatedGraph.composedSchemaVersionId,
Expand Down
4 changes: 4 additions & 0 deletions controlplane/src/core/env.schema.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -230,6 +230,10 @@ export const envVariables = z
* Admission Webhook
*/
AUTH_ADMISSION_JWT_SECRET: z.string(),
/**
* Prompt-to-Query
*/
PROMPT_TO_QUERY_SERVICE_ENDPOINT: z.string().url().optional(),
Comment on lines +233 to +236

@coderabbitai coderabbitai Bot Aug 5, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔒 Security & Privacy | 🟠 Major | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
set -eu

printf '%s\n' '--- references ---'
rg -n -C 5 --glob '!node_modules' --glob '!dist' --glob '!build' \
  'PROMPT_TO_QUERY_SERVICE_ENDPOINT|Prompt.?to.?Query|prompt.?to.?query' .

printf '%s\n' '--- candidate files ---'
git ls-files | rg -i 'prompt|query|env|docker|helm|kubernetes|compose|deployment|readme|config' | head -200

Repository: wundergraph/cosmo

Length of output: 19182

🏁 Script executed:

#!/bin/bash
set -eu

printf '%s\n' '--- PromptToQueryService structure ---'
ast-grep outline controlplane/src/core/services/PromptToQueryService.ts

printf '%s\n' '--- service implementation ---'
sed -n '1,190p' controlplane/src/core/services/PromptToQueryService.ts

printf '%s\n' '--- HTTP client and promptToQuery wiring ---'
rg -n -C 8 --glob '*.ts' --glob '*.tsx' \
  'promptToQuery|serviceAddress|httpClient|GenerateQuery|YokoService' controlplane/src

printf '%s\n' '--- configuration examples and deployment references ---'
rg -n -C 6 --glob '!node_modules' --glob '!dist' --glob '!build' \
  'PROMPT_TO_QUERY_SERVICE_ENDPOINT|prompt-to-query|Prompt To Query' \
  controlplane docker docker-compose.yml docker-compose.full.yml docker-compose.cosmo.yml \
  .github docs-website 2>/dev/null || true

Repository: wundergraph/cosmo

Length of output: 50374

🏁 Script executed:

#!/bin/bash
set -eu

printf '%s\n' '--- exact endpoint configuration references ---'
rg -n -C 8 --glob '!node_modules' --glob '!dist' --glob '!build' \
  'PROMPT_TO_QUERY_SERVICE_ENDPOINT' controlplane .github docs-website docker docker-compose.yml docker-compose.full.yml docker-compose.cosmo.yml

printf '%s\n' '--- PTQ method callers and request entry points ---'
rg -n -C 6 --glob '*.ts' --glob '*.tsx' \
  '\.generateQuery\(|generateQuery\(|\.indexSchema\(|indexSchema\(' controlplane/src

printf '%s\n' '--- deployment transport controls near controlplane configuration ---'
rg -n -C 5 --glob '*.yaml' --glob '*.yml' --glob '*.md' --glob '*.env.example' \
  'PROMPT_TO_QUERY|controlplane.*(url|host)|https://|http://' \
  .github docs-website controlplane docker docker-compose.yml docker-compose.full.yml docker-compose.cosmo.yml 2>/dev/null | head -300

Repository: wundergraph/cosmo

Length of output: 29384

Security Misconfiguration (CWE-319): Cleartext Transmission of Sensitive Information

Reachability: Internal

Reachability path 
● Entry
  controlplane/src/index.ts:83
  envVariables
│
▼
● Sink
  controlplane/src/core/env.schema.ts

Require HTTPS for PROMPT_TO_QUERY_SERVICE_ENDPOINT.

The PTQ client sends prompts and schema data through Axios. The schema accepts http://, which can expose this data on untrusted networks. Require HTTPS or enforce a private, authenticated network boundary.

🤖 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 `@controlplane/src/core/env.schema.ts` around lines 233 - 236, Update the
PROMPT_TO_QUERY_SERVICE_ENDPOINT schema definition in env.schema.ts to accept
only HTTPS URLs, replacing the current generic URL validation while preserving
its optional behavior.

})
.merge(sentryEnvVariables)
.refine((input) => {
Expand Down
1 change: 1 addition & 0 deletions controlplane/src/core/repositories/OrganizationRepository.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1438,6 +1438,7 @@ export class OrganizationRepository {
support: false,
'split-config-loading': false,
'login-method-restrictions': false,
'prompt-to-query': false,
};

for (const feature of features) {
Expand Down
1 change: 1 addition & 0 deletions controlplane/src/core/routes.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,7 @@ export interface RouterOptions {
stripeSecretKey?: string;
cdnBaseUrl: string;
lockAdapter: Redlock;
promptToQueryServiceAddress?: string;
}
const handlerOptions: Partial<ConnectRouterOptions> = {
maxTimeoutMs: 80_000,
Expand Down
Loading
Loading