Harden auth code normalization for pasted callback/code variants

Normalize --code inputs that are actually callback URLs or query fragments, strip wrapper punctuation/whitespace, and preserve inferred callback metadata/state for provider completion.

Also update Google skill command examples to quote callback/code arguments and add regression coverage for both callback-in-code and query-fragment inputs.

Co-Authored-By: GPT-5 Codex <codex@openai.com>

Author: dcramer

src/ash/capabilities/auth_normalization.py

@@ -30,11 +30,22 @@ def normalize_auth_completion(
     expected_state: str | None,
 ) -> NormalizedAuthCompletion:
     """Normalize callback URL / code inputs into one authorization code."""
-    normalized_code = _optional_text(code)
+    normalized_code: str | None = None
     normalized_callback_url = _optional_text(callback_url)
     callback_code: str | None = None
     callback_state: str | None = None
 
+    if code is not None:
+        (
+            normalized_code,
+            inferred_callback_url,
+            inferred_state,
+        ) = _normalize_code_input(code)
+        if normalized_callback_url is None and inferred_callback_url is not None:
+            normalized_callback_url = inferred_callback_url
+        if callback_state is None and inferred_state is not None:
+            callback_state = inferred_state
+
     if normalized_callback_url is not None:
         callback_code, callback_state = _parse_callback_url(normalized_callback_url)
 
@@ -96,3 +107,54 @@ def _optional_text(value: str | None) -> str | None:
         return None
     text = str(value).strip()
     return text or None
+
+
+def _normalize_code_input(code: str) -> tuple[str | None, str | None, str | None]:
+    text = _optional_text(code)
+    if text is None:
+        return None, None, None
+
+    cleaned = _trim_wrapper_punctuation(text)
+    compact = "".join(cleaned.split())
+
+    try:
+        code_from_url, state_from_url = _parse_callback_url(compact)
+    except AuthNormalizationError:
+        pass
+    else:
+        return code_from_url, compact, state_from_url
+
+    code_from_query, state_from_query = _extract_code_from_query_fragment(compact)
+    if code_from_query is not None:
+        return code_from_query, None, state_from_query
+
+    return compact, None, None
+
+
+def _extract_code_from_query_fragment(text: str) -> tuple[str | None, str | None]:
+    fragment = text
+    if "?" in fragment:
+        fragment = fragment.split("?", 1)[1]
+
+    if "code=" not in fragment:
+        return None, None
+
+    query = parse_qs(fragment, keep_blank_values=True)
+    code = _optional_text((query.get("code") or [None])[0])
+    state = _optional_text((query.get("state") or [None])[0])
+    return code, state
+
+
+def _trim_wrapper_punctuation(text: str) -> str:
+    cleaned = text.strip()
+    if len(cleaned) >= 2:
+        pairs = (("(", ")"), ("[", "]"), ("{", "}"), ("<", ">"), ("`", "`"))
+        changed = True
+        while changed and len(cleaned) >= 2:
+            changed = False
+            for left, right in pairs:
+                if cleaned.startswith(left) and cleaned.endswith(right):
+                    cleaned = cleaned[1:-1].strip()
+                    changed = True
+                    break
+    return cleaned.rstrip(".,;")

src/ash/integrations/skills/capabilities/google/SKILL.md

@@ -64,7 +64,7 @@ Run this step for each capability where `Authenticated: no`. If the user's reque
 Before prompting the user again, check whether the current task already contains a pasted Google callback URL (`http://localhost/?...code=...`) or a raw auth code. If yes:
 
 1. Run `ash-sb capability auth begin -c <capability>` first.
-2. Immediately run `ash-sb capability auth complete --flow-id <id> --callback-url <URL>` (or `--code <CODE>`).
+2. Immediately run `ash-sb capability auth complete --flow-id <id> --callback-url '<URL>'` (or `--code '<CODE>'`).
 3. Re-run `ash-sb capability list` and continue to operations if authenticated.
 
 Do not ask the user for another URL/code when one is already present in the task.
@@ -98,7 +98,7 @@ Started capability auth flow (flow_id=abc123)
 Check the `Flow type` in the output:
 
 - If `device_code`: show the `Auth URL` and `User code` from the output. Tell the user to open the URL and enter the code. Then proceed to step 2c to poll for completion.
-- If `authorization_code`: show the `Auth URL` from the output and ask the user to complete the Google consent screen and provide either the authorization code or the callback URL. Then use `ash-sb capability auth complete --flow-id <id> --code <CODE>` (or `--callback-url <URL>` when the user pastes the full callback URL). Do **not** pass `-c/--capability` to `auth complete`; that option is only valid for `auth begin`.
+- If `authorization_code`: show the `Auth URL` from the output and ask the user to complete the Google consent screen and provide either the authorization code or the callback URL. Then use `ash-sb capability auth complete --flow-id <id> --code '<CODE>'` (or `--callback-url '<URL>'` when the user pastes the full callback URL). Do **not** pass `-c/--capability` to `auth complete`; that option is only valid for `auth begin`.
 
 **2c. Poll for completion (device code flow)**
 

tests/test_capabilities.py

@@ -400,6 +400,63 @@ async def test_auth_complete_rejects_callback_state_mismatch() -> None:
     assert exc_info.value.code == "capability_auth_state_mismatch"
 
 
+@pytest.mark.asyncio
+async def test_auth_complete_accepts_callback_url_in_code_field() -> None:
+    manager = CapabilityManager(auth_flow_ttl_seconds=300)
+    provider = _RecordingProvider(namespace="gog")
+    await manager.register_provider(provider)
+    begin = await manager.auth_begin(
+        capability_id="gog.email",
+        user_id="user-1",
+        chat_type="private",
+        account_hint="work",
+    )
+
+    flow_id = str(begin["flow_id"])
+    manager._auth_flows[flow_id].expected_callback_state = "expected-state"
+    callback = "http://localhost/?state=expected-state&code=abc123&scope=mail"
+
+    result = await manager.auth_complete(
+        flow_id=flow_id,
+        user_id="user-1",
+        callback_url=None,
+        code=callback,
+    )
+
+    assert result["ok"] is True
+    completion = provider.complete_calls[0]["completion"]
+    assert completion.authorization_code == "abc123"
+    assert completion.raw_callback_url == callback
+
+
+@pytest.mark.asyncio
+async def test_auth_complete_accepts_code_query_fragment_in_code_field() -> None:
+    manager = CapabilityManager(auth_flow_ttl_seconds=300)
+    provider = _RecordingProvider(namespace="gog")
+    await manager.register_provider(provider)
+    begin = await manager.auth_begin(
+        capability_id="gog.email",
+        user_id="user-1",
+        chat_type="private",
+        account_hint="work",
+    )
+
+    flow_id = str(begin["flow_id"])
+    manager._auth_flows[flow_id].expected_callback_state = "expected-state"
+
+    result = await manager.auth_complete(
+        flow_id=flow_id,
+        user_id="user-1",
+        callback_url=None,
+        code="?state=expected-state&code=abc123",
+    )
+
+    assert result["ok"] is True
+    completion = provider.complete_calls[0]["completion"]
+    assert completion.authorization_code == "abc123"
+    assert completion.state == "expected-state"
+
+
 @pytest.mark.asyncio
 async def test_auth_begin_reuses_pending_flow_for_same_scope() -> None:
     manager = CapabilityManager(auth_flow_ttl_seconds=300)