Skip to content

Commit b86f0d6

Browse files
feat(oss): ship E2E fix proposal MVP (W1–W5 + thin W3/W6)
Close the OSS launch blockers from docs/oss-ux/E2E_FIX_PROPOSAL.md: - W2a/b: Postgres Brain scans all non-system schemas with qualified keys; coverage gate fails or NEEDS_ATTENTION instead of fake Complete 100% - W1: provisioner writes DEEPSQL_TOKEN_FILE, fail-loud provision, MCP auth probe + Agent boot banner, revoke cleans disk token - W4: Brain stage enum sync, stepper advances to Add context, jobs collapsed - W5: index enrichment + skip UNINDEXED_* on PK/UK/TRUE_KEY - W3 thin: /onboarding routed + redirect when no connections; payload fix - W6: DeepSQL title, generic Agent suggestions, dbType canonicalize Co-authored-by: Venkat SF <venkatesh.sakamuri@stayflexi.com>
1 parent 575f345 commit b86f0d6

26 files changed

Lines changed: 1067 additions & 371 deletions

backend/src/main/java/com/dbaagent/controller/AgentBridgeController.java

Lines changed: 33 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,10 +5,14 @@
55
import jakarta.servlet.http.Cookie;
66
import jakarta.servlet.http.HttpServletRequest;
77
import lombok.RequiredArgsConstructor;
8+
import org.slf4j.Logger;
9+
import org.slf4j.LoggerFactory;
810
import org.springframework.beans.factory.annotation.Value;
11+
import org.springframework.http.HttpStatus;
912
import org.springframework.http.ResponseEntity;
1013
import org.springframework.web.bind.annotation.*;
1114

15+
import java.util.HashMap;
1216
import java.util.Map;
1317

1418
/**
@@ -20,6 +24,8 @@
2024
@RequestMapping("/agent")
2125
@RequiredArgsConstructor
2226
public class AgentBridgeController {
27+
private static final Logger log = LoggerFactory.getLogger(AgentBridgeController.class);
28+
2329
private final AccessControlService accessControlService;
2430
private final AgentBridgeService agentBridgeService;
2531

@@ -33,8 +39,33 @@ public ResponseEntity<Map<String, Object>> session(
3339
String username = accessControlService.requireCurrentUsername();
3440
String token = extractToken(request);
3541
String connectionId = body == null ? null : asString(body.get("connectionId"));
36-
String profile = agentBridgeService.ensureProfile(username, token, connectionId);
37-
return ResponseEntity.ok(Map.of("profile", profile, "username", username));
42+
43+
AgentBridgeService.ProfileBootstrap bootstrap;
44+
try {
45+
bootstrap = agentBridgeService.ensureProfile(username, token, connectionId);
46+
} catch (AgentBridgeService.ProvisioningException e) {
47+
log.warn("Agent session bootstrap failed for {}: {}", username, e.getMessage());
48+
return ResponseEntity.status(HttpStatus.SERVICE_UNAVAILABLE).body(Map.of(
49+
"error", "agent_provisioning_failed",
50+
"message", "Could not provision the DeepSQL Agent for this user. "
51+
+ "Check that the agent provisioner is running and reachable."
52+
));
53+
}
54+
55+
// Boot health check: probe the exact token just provisioned against this
56+
// backend's own API before telling the UI it's safe to chat. Without
57+
// this, an expired/misrouted token surfaces only after several failed
58+
// tool calls deep into a conversation (W1 "fail loud, early").
59+
Map<String, Object> response = new HashMap<>();
60+
response.put("profile", bootstrap.profile());
61+
response.put("username", username);
62+
boolean mcpAuthOk = agentBridgeService.probeMcpAuth(bootstrap.token());
63+
response.put("mcpAuthOk", mcpAuthOk);
64+
if (!mcpAuthOk) {
65+
response.put("mcpAuthError", "The DeepSQL Agent could not authenticate against this API with its "
66+
+ "provisioned token. Reconnect or check the Agent runtime.");
67+
}
68+
return ResponseEntity.ok(response);
3869
}
3970

4071
private String extractToken(HttpServletRequest request) {

backend/src/main/java/com/dbaagent/model/InitStage.java

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,11 +11,16 @@ public enum InitStage {
1111
RAG_EMBEDDING,
1212
BRAIN_ANALYSIS,
1313
SEMANTIC_MODELING,
14+
/**
15+
* Schema coverage incomplete — Brain must not claim Complete 100%.
16+
* Surfaced when indexed base tables are far below live user tables (W2b).
17+
*/
18+
NEEDS_ATTENTION,
1419
COMPLETED,
1520
FAILED;
1621

1722
public boolean isTerminal() {
18-
return this == COMPLETED || this == FAILED;
23+
return this == COMPLETED || this == FAILED || this == NEEDS_ATTENTION;
1924
}
2025

2126
public InitStage next() {
@@ -29,7 +34,7 @@ public InitStage next() {
2934
case AI_DESCRIPTION -> RAG_EMBEDDING;
3035
case RAG_EMBEDDING -> BRAIN_ANALYSIS;
3136
case BRAIN_ANALYSIS -> SEMANTIC_MODELING;
32-
case SEMANTIC_MODELING, COMPLETED, FAILED -> null;
37+
case SEMANTIC_MODELING, COMPLETED, FAILED, NEEDS_ATTENTION -> null;
3338
};
3439
}
3540
}

backend/src/main/java/com/dbaagent/provider/postgres/PostgresIntrospectionProvider.java

Lines changed: 157 additions & 97 deletions
Large diffs are not rendered by default.

backend/src/main/java/com/dbaagent/service/AgentBridgeService.java

Lines changed: 159 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -26,8 +26,10 @@
2626
*
2727
* <p>The Java backend can't run the agent shell tooling itself (no agent runtime
2828
* in its image), so provisioning is delegated over the compose network to the
29-
* agent's internal, secret-gated provisioner endpoint. Failures never block the
30-
* tab — we still return the profile name.
29+
* agent's internal, secret-gated provisioner endpoint. The browser Agent-tab
30+
* boot path ({@link #ensureProfile}) is fail-loud — a configured-but-failing
31+
* provisioner throws rather than returning a profile with a stale disk token.
32+
* Headless channels ({@link #ensureProfileForUser}) stay best-effort.
3133
*/
3234
@Service
3335
public class AgentBridgeService {
@@ -91,24 +93,60 @@ public AgentBridgeService(McpTokenService mcpTokenService) {
9193
@Value("${security.session.refresh-days:7}")
9294
private long sessionWindowDays;
9395

96+
/**
97+
* Base URL for this same backend's own REST API, used only by
98+
* {@link #probeMcpAuth} to verify a freshly minted token actually works
99+
* before the Agent tab opens chat. Loopback by default — the probe never
100+
* needs to leave the box, so no external base-URL config is required.
101+
*/
102+
@Value("${agent.local-api-base-url:http://127.0.0.1:${server.port:8080}/api}")
103+
private String localApiBaseUrl;
104+
94105
public String profileFor(String username) {
95106
String safe = username.toLowerCase().replaceAll("[^a-z0-9]+", "-").replaceAll("(^-+|-+$)", "");
96107
return "u-" + safe;
97108
}
98109

110+
/**
111+
* Thrown when provisioning is configured (enabled + secret set) but the
112+
* provisioner call itself fails — non-2xx response or a connect/timeout
113+
* error. Callers must surface this rather than silently returning a
114+
* profile name that may point at a stale disk token (the W1 fix for
115+
* "Agent tab opens, tool calls 401 six steps later").
116+
*/
117+
public static class ProvisioningException extends RuntimeException {
118+
public ProvisioningException(String message) {
119+
super(message);
120+
}
121+
public ProvisioningException(String message, Throwable cause) {
122+
super(message, cause);
123+
}
124+
}
125+
126+
/** Result of {@link #ensureProfile}: the resolved profile plus the token actually provisioned with it. */
127+
public record ProfileBootstrap(String profile, String token) {}
128+
99129
/**
100130
* Ensure the user's agent profile exists and is bound to their current token;
101-
* returns the profile name. Best-effort — provisioning problems are logged but
102-
* never thrown (the Agent tab still opens, just without fresh per-user scope).
131+
* returns the profile name plus the token that was provisioned with it (so
132+
* the caller can probe that exact credential — see
133+
* {@link AgentBridgeService#probeMcpAuth}).
134+
*
135+
* <p>Fail-loud: when provisioning is configured (enabled + secret set), a
136+
* non-2xx or unreachable provisioner throws {@link ProvisioningException}
137+
* instead of returning a profile that may still carry a stale/expired disk
138+
* token. When provisioning is disabled or unconfigured, the tab still opens
139+
* against the shared default profile (documented, not silent) and the
140+
* returned token is the caller-supplied session token, unprovisioned.
103141
*/
104-
public String ensureProfile(String username, String authToken, String connectionId) {
142+
public ProfileBootstrap ensureProfile(String username, String authToken, String connectionId) {
105143
String profile = profileFor(username);
106144
if (!provisionEnabled) {
107-
return profile;
145+
return new ProfileBootstrap(profile, authToken);
108146
}
109147
if (provisionSecret == null || provisionSecret.isBlank()) {
110148
log.warn("agent.provision-secret is unset — skipping per-user provisioning for {}", username);
111-
return profile;
149+
return new ProfileBootstrap(profile, authToken);
112150
}
113151
// The agent profile must carry a credential that outlives a single chat
114152
// session. The user's session JWT lives only ~15 min and is coupled to a
@@ -122,16 +160,21 @@ public String ensureProfile(String username, String authToken, String connection
122160
agentToken = authToken == null ? "" : authToken;
123161
}
124162
callProvisioner(username, profile, agentToken, connectionId);
125-
return profile;
163+
return new ProfileBootstrap(profile, agentToken);
126164
}
127165

128166
/**
129-
* Headless variant for non-browser channels (Slack, etc.): provision the
130-
* user's profile with a dedicated CHANNEL token minted directly for the
131-
* DeepSQL user — no inbound session/cookie needed. The channel token has its
132-
* own name so the UI login/logout lifecycle (which extends/revokes the
133-
* {@code (auto)} token) never touches it; a web logout won't kill the user's
134-
* Slack agent access. Best-effort — never throws.
167+
* Headless variant for non-browser channels (Slack, dashboard generation,
168+
* the agent-chat turn API): provision the user's profile with a dedicated
169+
* CHANNEL token minted directly for the DeepSQL user — no inbound
170+
* session/cookie needed. The channel token has its own name so the UI
171+
* login/logout lifecycle (which extends/revokes the {@code (auto)} token)
172+
* never touches it; a web logout won't kill the user's Slack agent access.
173+
*
174+
* <p>Unlike {@link #ensureProfile} (the browser Agent-tab boot path, which
175+
* is fail-loud), this stays best-effort: a provisioner hiccup here must not
176+
* take down dashboard generation or a Slack turn over a transient network
177+
* blip. Provisioning failures are logged, not propagated.
135178
*/
136179
public String ensureProfileForUser(String username, String connectionId) {
137180
String profile = profileFor(username);
@@ -147,11 +190,24 @@ public String ensureProfileForUser(String username, String connectionId) {
147190
log.warn("Could not mint channel token for {} — skipping headless provisioning", username);
148191
return profile;
149192
}
150-
callProvisioner(username, profile, channelToken, connectionId);
193+
try {
194+
callProvisioner(username, profile, channelToken, connectionId);
195+
} catch (ProvisioningException e) {
196+
log.warn("Headless agent provisioning failed for {}: {}", username, e.getMessage());
197+
}
151198
return profile;
152199
}
153200

154-
/** POST the provision request to the agent container's internal provisioner. */
201+
/**
202+
* POST the provision request to the agent container's internal provisioner.
203+
*
204+
* <p>Throws {@link ProvisioningException} on a non-2xx response or any
205+
* connect/timeout/IO failure — the caller (both {@code ensureProfile}
206+
* variants) has already confirmed provisioning is enabled and configured, so
207+
* a failure here means the Agent tab is about to open against a profile the
208+
* provisioner never actually refreshed. That must block the tab, not log a
209+
* warning and proceed.
210+
*/
155211
private void callProvisioner(String username, String profile, String token, String connectionId) {
156212
try {
157213
String body = objectMapper.writeValueAsString(Map.of(
@@ -168,10 +224,91 @@ private void callProvisioner(String username, String profile, String token, Stri
168224
if (resp.statusCode() / 100 == 2) {
169225
log.info("Provisioned/refreshed agent profile {} for user {}", profile, username);
170226
} else {
171-
log.warn("Agent provisioning HTTP {} for user {}: {}", resp.statusCode(), username, resp.body());
227+
String message = "Agent provisioning HTTP " + resp.statusCode() + " for user " + username
228+
+ ": " + resp.body();
229+
log.warn(message);
230+
throw new ProvisioningException(message);
172231
}
232+
} catch (ProvisioningException e) {
233+
throw e;
173234
} catch (Exception e) {
174235
log.warn("Agent provisioning call failed for user {}: {}", username, e.getMessage());
236+
throw new ProvisioningException(
237+
"Agent provisioning call failed for user " + username + ": " + e.getMessage(), e);
238+
}
239+
}
240+
241+
/**
242+
* Derive the provisioner's revoke endpoint from its configured provision
243+
* URL ({@code .../provision} -> {@code .../revoke}). Returns null if the
244+
* configured URL doesn't follow that convention (best-effort only).
245+
*/
246+
private String revokeUrl() {
247+
if (provisionerUrl == null || !provisionerUrl.contains("/provision")) {
248+
return null;
249+
}
250+
return provisionerUrl.replace("/provision", "/revoke");
251+
}
252+
253+
/**
254+
* Best-effort POST to the provisioner's {@code /revoke} so the on-disk
255+
* token file / env fallback are cleared alongside the DB-side revoke.
256+
* Never throws — this runs after the DB token is already gone, so a
257+
* provisioner hiccup here must not fail the logout request itself.
258+
*/
259+
private void callProvisionerRevoke(String username) {
260+
if (!provisionEnabled) {
261+
return;
262+
}
263+
String url = revokeUrl();
264+
if (url == null || provisionSecret == null || provisionSecret.isBlank()) {
265+
return;
266+
}
267+
try {
268+
String body = objectMapper.writeValueAsString(Map.of("user", username));
269+
HttpRequest req = HttpRequest.newBuilder(URI.create(url))
270+
.header("Content-Type", "application/json")
271+
.header("X-Provision-Secret", provisionSecret)
272+
.timeout(Duration.ofSeconds(10))
273+
.POST(HttpRequest.BodyPublishers.ofString(body))
274+
.build();
275+
HttpResponse<String> resp = http.send(req, HttpResponse.BodyHandlers.ofString());
276+
if (resp.statusCode() / 100 == 2) {
277+
log.info("Revoked on-disk agent token for user {}", username);
278+
} else {
279+
log.warn("Agent token revoke HTTP {} for user {}: {}", resp.statusCode(), username, resp.body());
280+
}
281+
} catch (Exception e) {
282+
log.warn("Agent token revoke call failed for user {}: {}", username, e.getMessage());
283+
}
284+
}
285+
286+
/**
287+
* Probe whether the just-minted/refreshed MCP token can actually reach the
288+
* DeepSQL API — the health check the Agent tab boot depends on to decide
289+
* whether to show the chat composer or a blocking "Agent cannot reach
290+
* DeepSQL (auth)" banner. GETs {@code /connections} with the token as a
291+
* bearer credential against the local backend (loopback — this call never
292+
* leaves the box, so no external base-URL config is needed).
293+
*
294+
* @return true if the API accepted the token (2xx), false on any
295+
* non-2xx/auth failure or network error.
296+
*/
297+
public boolean probeMcpAuth(String token) {
298+
if (token == null || token.isBlank()) {
299+
return false;
300+
}
301+
try {
302+
HttpRequest req = HttpRequest.newBuilder(URI.create(localApiBaseUrl + "/connections"))
303+
.header("Authorization", "Bearer " + token)
304+
.timeout(Duration.ofSeconds(5))
305+
.GET()
306+
.build();
307+
HttpResponse<String> resp = http.send(req, HttpResponse.BodyHandlers.ofString());
308+
return resp.statusCode() / 100 == 2;
309+
} catch (Exception e) {
310+
log.warn("MCP auth probe failed: {}", e.getMessage());
311+
return false;
175312
}
176313
}
177314

@@ -258,5 +395,10 @@ public void revokeAgentTokens(String username) {
258395
} catch (Exception e) {
259396
log.warn("Could not revoke agent token(s) for {}: {}", username, e.getMessage());
260397
}
398+
// DB-side revoke only kills future auth checks — the plaintext token can
399+
// still live on disk in the profile's .env / token file / MCP server env
400+
// until the provisioner overwrites it. Clear that copy too so a revoked
401+
// token can't keep the agent working. Best-effort: never blocks logout.
402+
callProvisionerRevoke(username);
261403
}
262404
}

backend/src/main/java/com/dbaagent/service/CredentialService.java

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import com.dbaagent.model.ConnectionRequest;
44
import com.dbaagent.model.DatabaseConnection;
5+
import com.dbaagent.provider.DatabaseProviderRegistry;
56
import com.dbaagent.repository.CredentialRepository;
67
import com.dbaagent.security.EncryptionService;
78
import com.dbaagent.service.security.ConnectionAccessService;
@@ -26,13 +27,18 @@ public class CredentialService {
2627
private final EncryptionService encryptionService;
2728
private final ConnectionAccessService connectionAccessService;
2829
private final TelemetryClient telemetryClient;
30+
private final DatabaseProviderRegistry providerRegistry;
2931

3032
@Transactional
3133
public DatabaseConnection saveConnection(ConnectionRequest request, String ownerUsername) {
3234
DatabaseConnection connection = new DatabaseConnection();
3335
connection.setId(UUID.randomUUID().toString());
3436
connection.setConnectionName(request.getConnectionName());
35-
connection.setDbType(request.getDbType());
37+
// Canonicalize through the provider registry ("postgresql" -> "postgres", etc.) so
38+
// every downstream consumer that switches on dbType (DatabaseProviderRegistry.getDialect,
39+
// frontend badges, telemetry) sees one spelling per dialect regardless of which alias
40+
// the caller (onboarding wizard, API client, import) happened to send.
41+
connection.setDbType(providerRegistry.getCanonicalName(request.getDbType()));
3642
connection.setOwnerUsername(ownerUsername);
3743
connection.setCreatedAt(LocalDateTime.now());
3844
connection.setLastUsed(LocalDateTime.now());
@@ -377,7 +383,7 @@ public DatabaseConnection updateConnection(String connectionId, ConnectionReques
377383

378384
// Update non-encrypted fields
379385
connection.setConnectionName(request.getConnectionName());
380-
connection.setDbType(request.getDbType());
386+
connection.setDbType(providerRegistry.getCanonicalName(request.getDbType()));
381387
connection.setLastUsed(LocalDateTime.now());
382388

383389
// Update and re-encrypt sensitive fields

0 commit comments

Comments
 (0)