Skip to content

Commit 19d9562

Browse files
henrypark133claude
henrypark133
and
claude
authored
feat(extensions): unify auth and configure into single entrypoint (#677)
* feat(extensions): unify auth and configure into single entrypoint Refactors the extension lifecycle to eliminate the divergence between chat and gateway paths that caused Telegram setup via chat to fail (missing webhook secret auto-generation, no token validation). Key changes: - Rename save_setup_secrets() → configure(): single entrypoint for providing secrets to any extension (WasmChannel, WasmTool, MCP). Validates, stores, auto-generates, and activates. - Add configure_token(): convenience wrapper for single-token callers (chat auth card, WebSocket, agent auth mode). - Refactor auth() to pure status check: remove token parameter, delete token-storing branches from auth_mcp/auth_wasm_tool, rename auth_wasm_channel → auth_wasm_channel_status. - Add ConfigureResult/MissingSecret types for structured responses. - Replace hardcoded Telegram token validation with generic validation_endpoint from capabilities.json. - Update all callers (9 files) to use the new interface. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: use ValidationFailed error variant instead of string matching Replace brittle msg.contains("Invalid token") checks with a proper ExtensionError::ValidationFailed variant. configure() now returns this variant for token validation failures, and callers match on it directly instead of parsing error message strings. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: address review — SSRF protection, error typing, missing-secret selection, WS auth 1. SSRF: call validate_fetch_url() before validation_endpoint HTTP request 2. Transport errors map to ExtensionError::Other (not ValidationFailed) 3. configure_token() picks first *missing* secret, not first non-optional 4. WebSocket error path re-emits AuthRequired on ValidationFailed Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * test: add regression tests for extension lifecycle refactoring - test_configure_token_picks_first_missing_secret: verifies multi-secret channels can be configured one secret at a time (commit ce106f4) - test_auth_is_read_only_for_wasm_channel: verifies auth() has no side effects and doesn't store secrets (commit 47f8eb6) - test_validation_failed_is_distinct_error_variant: verifies the typed error variant can be pattern-matched (commit a318161) Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * fix: address review comments — activation dispatch, dead code, caps consolidation - Fix configure() fallthrough bug: dispatch activation by ExtensionKind instead of unconditionally calling activate_wasm_channel() for all non-WasmTool types (MCP servers and channel relays now use their correct activation methods) - Remove dead MissingSecret struct and missing_secrets field (never populated, flagged by reviewer) - Consolidate capabilities file parsing in configure(): parse once and reuse for allowed names, validation_endpoint, and auto-generation - Fix auth() doc comment: note MCP OAuth side effects - Fix stale save_setup_secrets reference in server.rs comment - Add regression test for activation dispatch bug Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent 28a22f2 commit 19d9562

8 files changed

Lines changed: 554 additions & 367 deletions

File tree

‎channels-src/telegram/telegram.capabilities.json‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,8 @@
2020
"optional": false
2121
}
2222
],
23-
"setup_url": "https://t.me/BotFather"
23+
"setup_url": "https://t.me/BotFather",
24+
"validation_endpoint": "https://api.telegram.org/bot{telegram_bot_token}/getMe"
2425
},
2526
"capabilities": {
2627
"http": {

‎src/agent/thread_ops.rs‎

Lines changed: 37 additions & 81 deletions
Original file line numberDiff line numberDiff line change
@@ -1355,100 +1355,56 @@ impl Agent {
13551355
None => return Ok(Some("Extension manager not available.".to_string())),
13561356
};
13571357

1358-
match ext_mgr.auth(&pending.extension_name, Some(token)).await {
1359-
Ok(result) if result.is_authenticated() => {
1358+
match ext_mgr
1359+
.configure_token(&pending.extension_name, token)
1360+
.await
1361+
{
1362+
Ok(result) => {
13601363
tracing::info!(
1361-
"Extension '{}' authenticated via auth mode",
1362-
pending.extension_name
1364+
"Extension '{}' configured via auth mode: {}",
1365+
pending.extension_name,
1366+
result.message
13631367
);
1364-
1365-
// Auto-activate so tools are available immediately after auth
1366-
match ext_mgr.activate(&pending.extension_name).await {
1367-
Ok(activate_result) => {
1368-
let tool_count = activate_result.tools_loaded.len();
1369-
let tool_list = if activate_result.tools_loaded.is_empty() {
1370-
String::new()
1371-
} else {
1372-
format!("\n\nTools: {}", activate_result.tools_loaded.join(", "))
1373-
};
1374-
let msg = format!(
1375-
"{} authenticated and activated ({} tools loaded).{}",
1376-
pending.extension_name, tool_count, tool_list
1377-
);
1378-
let _ = self
1379-
.channels
1380-
.send_status(
1381-
&message.channel,
1382-
StatusUpdate::AuthCompleted {
1383-
extension_name: pending.extension_name.clone(),
1384-
success: true,
1385-
message: msg.clone(),
1386-
},
1387-
&message.metadata,
1388-
)
1389-
.await;
1390-
Ok(Some(msg))
1391-
}
1392-
Err(e) => {
1393-
tracing::warn!(
1394-
"Extension '{}' authenticated but activation failed: {}",
1395-
pending.extension_name,
1396-
e
1397-
);
1398-
let msg = format!(
1399-
"{} authenticated successfully, but activation failed: {}. \
1400-
Try activating manually.",
1401-
pending.extension_name, e
1402-
);
1403-
let _ = self
1404-
.channels
1405-
.send_status(
1406-
&message.channel,
1407-
StatusUpdate::AuthCompleted {
1408-
extension_name: pending.extension_name.clone(),
1409-
success: true,
1410-
message: msg.clone(),
1411-
},
1412-
&message.metadata,
1413-
)
1414-
.await;
1415-
Ok(Some(msg))
1416-
}
1417-
}
1418-
}
1419-
Ok(result) => {
1420-
// Invalid token, re-enter auth mode
1421-
{
1422-
let mut sess = session.lock().await;
1423-
if let Some(thread) = sess.threads.get_mut(&thread_id) {
1424-
thread.enter_auth_mode(pending.extension_name.clone());
1425-
}
1426-
}
1427-
let msg = result
1428-
.instructions()
1429-
.map(String::from)
1430-
.unwrap_or_else(|| "Invalid token. Please try again.".to_string());
1431-
// Re-emit AuthRequired so web UI re-shows the card
14321368
let _ = self
14331369
.channels
14341370
.send_status(
14351371
&message.channel,
1436-
StatusUpdate::AuthRequired {
1372+
StatusUpdate::AuthCompleted {
14371373
extension_name: pending.extension_name.clone(),
1438-
instructions: Some(msg.clone()),
1439-
auth_url: result.auth_url().map(String::from),
1440-
setup_url: result.setup_url().map(String::from),
1374+
success: true,
1375+
message: result.message.clone(),
14411376
},
14421377
&message.metadata,
14431378
)
14441379
.await;
1445-
Ok(Some(msg))
1380+
Ok(Some(result.message))
14461381
}
14471382
Err(e) => {
1448-
let msg = format!(
1449-
"Authentication failed for {}: {}",
1450-
pending.extension_name, e
1451-
);
1383+
let msg = e.to_string();
1384+
// Token validation errors: re-enter auth mode and re-prompt
1385+
if matches!(e, crate::extensions::ExtensionError::ValidationFailed(_)) {
1386+
{
1387+
let mut sess = session.lock().await;
1388+
if let Some(thread) = sess.threads.get_mut(&thread_id) {
1389+
thread.enter_auth_mode(pending.extension_name.clone());
1390+
}
1391+
}
1392+
let _ = self
1393+
.channels
1394+
.send_status(
1395+
&message.channel,
1396+
StatusUpdate::AuthRequired {
1397+
extension_name: pending.extension_name.clone(),
1398+
instructions: Some(msg.clone()),
1399+
auth_url: None,
1400+
setup_url: None,
1401+
},
1402+
&message.metadata,
1403+
)
1404+
.await;
1405+
return Ok(Some(msg));
1406+
}
1407+
// Infrastructure errors
14521408
let _ = self
14531409
.channels
14541410
.send_status(

‎src/channels/web/handlers/chat.rs‎

Lines changed: 24 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -145,49 +145,33 @@ pub async fn chat_auth_token_handler(
145145
"Extension manager not available".to_string(),
146146
))?;
147147

148-
let result = ext_mgr
149-
.auth(&req.extension_name, Some(&req.token))
148+
match ext_mgr
149+
.configure_token(&req.extension_name, &req.token)
150150
.await
151-
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
152-
153-
if result.is_authenticated() {
154-
// Auto-activate so tools are available immediately
155-
let msg = match ext_mgr.activate(&req.extension_name).await {
156-
Ok(r) => format!(
157-
"{} authenticated ({} tools loaded)",
158-
req.extension_name,
159-
r.tools_loaded.len()
160-
),
161-
Err(e) => format!(
162-
"{} authenticated but activation failed: {}",
163-
req.extension_name, e
164-
),
165-
};
166-
167-
// Clear auth mode on the active thread
168-
clear_auth_mode(&state).await;
151+
{
152+
Ok(result) => {
153+
clear_auth_mode(&state).await;
169154

170-
state.sse.broadcast(SseEvent::AuthCompleted {
171-
extension_name: req.extension_name,
172-
success: true,
173-
message: msg.clone(),
174-
});
155+
state.sse.broadcast(SseEvent::AuthCompleted {
156+
extension_name: req.extension_name.clone(),
157+
success: true,
158+
message: result.message.clone(),
159+
});
175160

176-
Ok(Json(ActionResponse::ok(msg)))
177-
} else {
178-
// Re-emit auth_required for retry
179-
state.sse.broadcast(SseEvent::AuthRequired {
180-
extension_name: req.extension_name.clone(),
181-
instructions: result.instructions().map(String::from),
182-
auth_url: result.auth_url().map(String::from),
183-
setup_url: result.setup_url().map(String::from),
184-
});
185-
Ok(Json(ActionResponse::fail(
186-
result
187-
.instructions()
188-
.map(String::from)
189-
.unwrap_or_else(|| "Invalid token".to_string()),
190-
)))
161+
Ok(Json(ActionResponse::ok(result.message)))
162+
}
163+
Err(e) => {
164+
let msg = e.to_string();
165+
if matches!(e, crate::extensions::ExtensionError::ValidationFailed(_)) {
166+
state.sse.broadcast(SseEvent::AuthRequired {
167+
extension_name: req.extension_name.clone(),
168+
instructions: Some(msg.clone()),
169+
auth_url: None,
170+
setup_url: None,
171+
});
172+
}
173+
Ok(Json(ActionResponse::fail(msg)))
174+
}
191175
}
192176
}
193177

‎src/channels/web/server.rs‎

Lines changed: 31 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -1062,49 +1062,35 @@ async fn chat_auth_token_handler(
10621062
"Extension manager not available".to_string(),
10631063
))?;
10641064

1065-
let result = ext_mgr
1066-
.auth(&req.extension_name, Some(&req.token))
1065+
match ext_mgr
1066+
.configure_token(&req.extension_name, &req.token)
10671067
.await
1068-
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
1069-
1070-
if result.is_authenticated() {
1071-
// Auto-activate so tools are available immediately
1072-
let msg = match ext_mgr.activate(&req.extension_name).await {
1073-
Ok(r) => format!(
1074-
"{} authenticated ({} tools loaded)",
1075-
req.extension_name,
1076-
r.tools_loaded.len()
1077-
),
1078-
Err(e) => format!(
1079-
"{} authenticated but activation failed: {}",
1080-
req.extension_name, e
1081-
),
1082-
};
1083-
1084-
// Clear auth mode on the active thread
1085-
clear_auth_mode(&state).await;
1068+
{
1069+
Ok(result) => {
1070+
// Clear auth mode on the active thread
1071+
clear_auth_mode(&state).await;
10861072

1087-
state.sse.broadcast(SseEvent::AuthCompleted {
1088-
extension_name: req.extension_name,
1089-
success: true,
1090-
message: msg.clone(),
1091-
});
1073+
state.sse.broadcast(SseEvent::AuthCompleted {
1074+
extension_name: req.extension_name.clone(),
1075+
success: true,
1076+
message: result.message.clone(),
1077+
});
10921078

1093-
Ok(Json(ActionResponse::ok(msg)))
1094-
} else {
1095-
// Re-emit auth_required for retry
1096-
state.sse.broadcast(SseEvent::AuthRequired {
1097-
extension_name: req.extension_name.clone(),
1098-
instructions: result.instructions().map(String::from),
1099-
auth_url: result.auth_url().map(String::from),
1100-
setup_url: result.setup_url().map(String::from),
1101-
});
1102-
Ok(Json(ActionResponse::fail(
1103-
result
1104-
.instructions()
1105-
.map(String::from)
1106-
.unwrap_or_else(|| "Invalid token".to_string()),
1107-
)))
1079+
Ok(Json(ActionResponse::ok(result.message)))
1080+
}
1081+
Err(e) => {
1082+
let msg = e.to_string();
1083+
// Re-emit auth_required for retry on validation errors
1084+
if matches!(e, crate::extensions::ExtensionError::ValidationFailed(_)) {
1085+
state.sse.broadcast(SseEvent::AuthRequired {
1086+
extension_name: req.extension_name.clone(),
1087+
instructions: Some(msg.clone()),
1088+
auth_url: None,
1089+
setup_url: None,
1090+
});
1091+
}
1092+
Ok(Json(ActionResponse::fail(msg)))
1093+
}
11081094
}
11091095
}
11101096

@@ -1853,7 +1839,7 @@ async fn extensions_install_handler(
18531839
// expansion and for first-time auth when credentials are already
18541840
// configured (e.g., built-in providers). We only surface an auth_url
18551841
// when the extension reports it is awaiting authorization.
1856-
match ext_mgr.auth(&req.name, None).await {
1842+
match ext_mgr.auth(&req.name).await {
18571843
Ok(auth_result) if auth_result.auth_url().is_some() => {
18581844
// Scope expansion or initial OAuth: user needs to authorize
18591845
resp.auth_url = auth_result.auth_url().map(String::from);
@@ -1882,9 +1868,9 @@ async fn extensions_activate_handler(
18821868
// Activation loaded the WASM module. Check if the tool needs
18831869
// OAuth scope expansion (e.g., adding google-docs when gmail
18841870
// already has a token but missing the documents scope).
1885-
// Initial OAuth setup is triggered via save_setup_secrets.
1871+
// Initial OAuth setup is triggered via configure.
18861872
let mut resp = ActionResponse::ok(result.message);
1887-
if let Ok(auth_result) = ext_mgr.auth(&name, None).await
1873+
if let Ok(auth_result) = ext_mgr.auth(&name).await
18881874
&& auth_result.auth_url().is_some()
18891875
{
18901876
resp.auth_url = auth_result.auth_url().map(String::from);
@@ -1902,7 +1888,7 @@ async fn extensions_activate_handler(
19021888
}
19031889

19041890
// Activation failed due to auth; try authenticating first.
1905-
match ext_mgr.auth(&name, None).await {
1891+
match ext_mgr.auth(&name).await {
19061892
Ok(auth_result) if auth_result.is_authenticated() => {
19071893
// Auth succeeded, retry activation.
19081894
match ext_mgr.activate(&name).await {
@@ -2109,7 +2095,7 @@ async fn extensions_setup_submit_handler(
21092095
"Extension manager not available (secrets store required)".to_string(),
21102096
))?;
21112097

2112-
match ext_mgr.save_setup_secrets(&name, &req.secrets).await {
2098+
match ext_mgr.configure(&name, &req.secrets).await {
21132099
Ok(result) => {
21142100
// Broadcast auth_completed so the chat UI can dismiss any in-progress
21152101
// auth card or setup modal that was triggered by tool_auth/tool_activate.

‎src/channels/web/ws.rs‎

Lines changed: 15 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -255,43 +255,31 @@ async fn handle_client_message(
255255
token,
256256
} => {
257257
if let Some(ref ext_mgr) = state.extension_manager {
258-
match ext_mgr.auth(&extension_name, Some(&token)).await {
259-
Ok(result) if result.is_authenticated() => {
260-
let msg = match ext_mgr.activate(&extension_name).await {
261-
Ok(r) => format!(
262-
"{} authenticated ({} tools loaded)",
263-
extension_name,
264-
r.tools_loaded.len()
265-
),
266-
Err(e) => format!(
267-
"{} authenticated but activation failed: {}",
268-
extension_name, e
269-
),
270-
};
258+
match ext_mgr.configure_token(&extension_name, &token).await {
259+
Ok(result) => {
271260
crate::channels::web::server::clear_auth_mode(state).await;
272261
state
273262
.sse
274263
.broadcast(crate::channels::web::types::SseEvent::AuthCompleted {
275264
extension_name,
276265
success: true,
277-
message: msg,
278-
});
279-
}
280-
Ok(result) => {
281-
state
282-
.sse
283-
.broadcast(crate::channels::web::types::SseEvent::AuthRequired {
284-
extension_name,
285-
instructions: result.instructions().map(String::from),
286-
auth_url: result.auth_url().map(String::from),
287-
setup_url: result.setup_url().map(String::from),
266+
message: result.message,
288267
});
289268
}
290269
Err(e) => {
270+
let msg = format!("Auth failed: {}", e);
271+
if matches!(e, crate::extensions::ExtensionError::ValidationFailed(_)) {
272+
state.sse.broadcast(
273+
crate::channels::web::types::SseEvent::AuthRequired {
274+
extension_name: extension_name.clone(),
275+
instructions: Some(msg.clone()),
276+
auth_url: None,
277+
setup_url: None,
278+
},
279+
);
280+
}
291281
let _ = direct_tx
292-
.send(WsServerMessage::Error {
293-
message: format!("Auth failed: {}", e),
294-
})
282+
.send(WsServerMessage::Error { message: msg })
295283
.await;
296284
}
297285
}

0 commit comments

Comments
 (0)