Route only explicit Codex and Claude invocations through smart routing by lilly-luo · Pull Request #502 · databricks/ucode

lilly-luo · 2026-09-04T22:41:50Z

https://databricks.atlassian.net/browse/AIGTWY-4551

explanation

currently, we are invoking smart routing for /every/ codex and claude command. in reality, we only want to invoke smart routing in 2 cases:

claude / codex empty - this opens a tui and we route on the first prompt
claude -- 'hi' / codex -- 'hi' - this opens a tui with 'hi' as the first prompt.;

claude --model xx or codex --model yy should not invoke smart routing.

this PR makes 3 changes so that ucode codex upgrade and ucode codex app-server will work:

makes it so that we detect whether there's a prompt ('hi') or if there's no subcommand. if 1 of those is true, we invoke smart routing. otherwise, don't. this ensures things like ucode codex update and ucode codex app-server continue working without having to maintain a hardcoded subcommand list.
after (1), we'd get this error on ucode codex update. This is because we, by default, pass the ucode config as argument --profile to all codex commands. however, only codex and codex mcp take --profile. instead, the other commands take --config. so i change codex to pass in --config instead

  Error: --profile only applies to runtime commands and `codex mcp`: `codex`, `codex exec`, `codex review`, `codex resume`, `codex queue`, `codex
  archive`, `codex delete`, `codex unarchive`, `codex fork`, `codex mcp`, `codex sandbox`, and `codex debug prompt-input`.
  ! ucode's `--profile` isn't accepted here (error above). Retrying without it: Codex will resolve /home/lilly.luo/.codex/config.toml and any
  OS-managed settings instead of the ucode profile.

even after (2), ucode codex upgrade would return the error below. this is because when we pass --config instead of --profile to codex, we need to serialize the toml file. the serialization logic had to be updated as a result of tomlkit upgrading to 15.0.

finally after 1, 2, and 3 - ucode codex upgrade and ucode codex app-server work.

  ValueError: Inline tables cannot contain a table

test

ENABLE_SMART_ROUTING_V2=1 SMART_ROUTER_NAME=task_v2 uv run ucode claude
ENABLE_SMART_ROUTING_V2=1 SMART_ROUTER_NAME=task_v2 uv run ucode codex
ENABLE_SMART_ROUTING_V2=1 SMART_ROUTER_NAME=task_v2 uv run ucode claude -- 'hi'
ENABLE_SMART_ROUTING_V2=1 SMART_ROUTER_NAME=task_v2 uv run ucode codex -- 'hi'
ENABLE_SMART_ROUTING_V2=1 SMART_ROUTER_NAME=task_v2 uv run ucode claude --model system.ai.glm-5-2
ENABLE_SMART_ROUTING_V2=1 SMART_ROUTER_NAME=task_v2 uv run ucode codex --model system.ai.glm-5-2
ENABLE_SMART_ROUTING_V2=1 SMART_ROUTER_NAME=task_v2 uv run ucode codex update
ENABLE_SMART_ROUTING_V2=1 SMART_ROUTER_NAME=task_v2 uv run ucode claude upgrade
ENABLE_SMART_ROUTING_V2=1 SMART_ROUTER_NAME=task_v2 uv run ucode codex app-server

the below cmds should be unchanged, i tested to ensure no regression

ENABLE_SMART_ROUTING_V2=0 uv run ucode claude
ENABLE_SMART_ROUTING_V2=0 uv run ucode codex
ENABLE_SMART_ROUTING_V2=0 uv run ucode claude --model system.ai.glm-5-2
ENABLE_SMART_ROUTING_V2=0 uv run ucode codex --model system.ai.glm-5-2
ENABLE_SMART_ROUTING_V2=0 uv run ucode claude -- 'hi'
ENABLE_SMART_ROUTING_V2=0 uv run ucode codex -- 'hi'
ENABLE_SMART_ROUTING_V2=0 uv run ucode codex update
ENABLE_SMART_ROUTING_V2=0 uv run ucode claude upgrade
ENABLE_SMART_ROUTING_V2=0 uv run ucode codex app-server

before (confirming i repro'd the issue) :

# version of ucode currently in isaac 
$ uv tool install \
  --default-index https://pypi-proxy.cloud.databricks.com/simple \
  'ucode==0.1.0+f7b4b97'

$ ENABLE_SMART_ROUTING_V2=1 ucode codex update    

╭────────────────────────────────────╮
│ Launching Codex with Unity Gateway │
╰────────────────────────────────────╯
✔ Starting Codex
Error: `--remote ws://127.0.0.1:40091` is only supported for interactive TUI commands, not `codex update`

$ ENABLE_SMART_ROUTING_V2=1 ucode codex app-server

╭────────────────────────────────────╮
│ Launching Codex with Unity Gateway │
╰────────────────────────────────────╯
✔ Starting Codex
Error: `--remote ws://127.0.0.1:32895` is only supported for interactive TUI commands, not `codex app-server`

after (it's fixed)

$ ENABLE_SMART_ROUTING_V2=1 uv run ucode codex update

╭────────────────────────────────────╮
│ Launching Codex with Unity Gateway │
╰────────────────────────────────────╯
✔ Starting Codex

Updating Codex via `bun install -g @openai/codex`...
bun add v1.3.14 (0d9b296a)

installed @openai/codex@0.153.4 with binaries:
 - codex

[1239.00ms] done

🎉 Update ran successfully! Please restart Codex.

$ ENABLE_SMART_ROUTING_V2=1 uv run ucode codex app-server (the 404 is unrelated and still happens with normal codex) 

╭────────────────────────────────────╮
│ Launching Codex with Unity Gateway │
╰────────────────────────────────────╯
✔ Starting Codex
2026-09-04T23:54:14.795144Z ERROR codex_models_manager::manager: failed to refresh available models: unexpected status 404 Not Found: {"error_code":"ENDPOINT_NOT_FOUND","message":"codex/v1/models is not enabled for this workspace."}, url: https://dbc-a5d4177a-49dc.cloud.databricks.com/ai-gateway/codex/v1/models?client_version=0.153.4, request id: fde50f25-d7db-4145-a011-9a49e892670b

…s-codex

lilly-luo · 2026-09-05T00:06:48Z

-        version_text = agent_version(binary)
-        parsed_version = _parse_version(version_text)
-        if parsed_version is not None and parsed_version < MINIMUM_ROUTING_CODEX_VERSION:
-            raise RuntimeError(
-                "Codex smart routing requires Codex "
-                f"{MINIMUM_ROUTING_CODEX_VERSION_TEXT} or newer; found {version_text}."
-            )
-
-        def _app_server_start_model() -> str:
-            managed_model = default_model(state)
-            if managed_model:
-                return managed_model
-            models = routing_models(state)
-            if models:
-                return codex_model_id(models[0])
-            return APP_SERVER_SMART_ROUTING_STARTING_MODEL
-
-        smart_routing_v2.launch_codex(
-            state,
-            tool_args,
-            binary=binary,
-            start_model=_app_server_start_model(),
-            render_overlay=render_overlay,
-        )


this is all moved to _launch_smart_routing and is identical

lilly-luo · 2026-09-05T00:45:43Z

-def _uses_interactive_tui(tool_args: list[str]) -> bool:
-    if any(arg in CLAUDE_NONINTERACTIVE_FLAGS for arg in tool_args):
-        return False
-
-    index = 0
-    while index < len(tool_args):
-        arg = tool_args[index]
-        if arg == "--":
-            return True
-        if arg in CLAUDE_VALUE_OPTIONS:
-            index += 2
-            continue
-        if arg in CLAUDE_OPTIONAL_VALUE_OPTIONS:
-            if index + 1 < len(tool_args) and not tool_args[index + 1].startswith("-"):
-                index += 2
-            else:
-                index += 1
-            continue
-        if arg.startswith("-"):
-            index += 1
-            continue
-        # Claude accepts an initial prompt positionally and still opens the TUI.
-        # Keep prompts inside the V2 PTY while bypassing utility subcommands.
-        return arg not in CLAUDE_SUBCOMMANDS
-    return True
-
-


this was very much not needed. now, we just look and see if claude either (1) has no subcommand or (2) has a -- . if (1) or (2) is true, we smart route

lilly-luo · 2026-09-05T00:46:25Z

-    if tool_args[:1] == ["app"]:
-        # `codex app` rejects --profile. Pass the ucode profile as --config
-        # overrides instead, preserving its Databricks provider and auth
-        # settings without changing the user's base config.toml.
-        profile_doc = read_toml_safe(CODEX_CONFIG_PATH)
-        if not profile_doc:
-            raise RuntimeError(
-                f"Cannot launch Codex app with the ucode profile because {CODEX_CONFIG_PATH} "
-                "is missing or empty. Run `ucode configure --agents codex` first."
-            )
-        config_args = codex_config_args(profile_doc)
-        exec_or_spawn([binary, "app", *config_args, *tool_args[1:]])
-        return  # unreachable in production (exec replaces the process)
-    # Run codex with --profile first — the TUI and runtime subcommands
-    # (exec/resume/mcp/...) keep ucode's Databricks routing, including any added
-    # by future codex versions. codex rejects the global --profile on
-    # server-family subcommands (app-server, mcp-server, ...), which are
-    # caller-configured anyway (e.g. omnigent runs `codex app-server` with its
-    # own CODEX_HOME); on that rejection we relaunch without --profile.
-    #
-    # The retry is gated on the attempt failing *fast*: the rejection is a
-    # parse-time error (~0.15s), whereas a session that actually starts can only


this was all very much not needed. we can just pass in --config for everything and do the parsing from toml -> config

lilly-luo · 2026-09-05T00:47:54Z

this file introduces parsing of the tomlkit item in case it is nested.

lilly-luo force-pushed the stack/lilly/only-do-smart-routing-for-2-cases-codex branch from 6dd7576 to dcda6c4 Compare September 4, 2026 22:58

only do smart routing for 2 cases: no subcommand or the --

586eeef

lilly-luo force-pushed the stack/lilly/only-do-smart-routing-for-2-cases-codex branch from dcda6c4 to 586eeef Compare September 4, 2026 23:10

lilly-luo changed the title ~~Route only explicit Codex TUI invocations through smart routing~~ only route codex / claude (no prompt) and "codex -- '<>'" and "claude -- '<>'" to smart router Sep 4, 2026

Merge branch 'main' into stack/lilly/only-do-smart-routing-for-2-case…

ef9dfc9

…s-codex

lilly-luo changed the title ~~only route codex / claude (no prompt) and "codex -- '<>'" and "claude -- '<>'" to smart router~~ only route codex / claude (no prompt) and codex -- '<>' and claude -- '<>' to smart router Sep 4, 2026

lilly-luo force-pushed the stack/lilly/only-do-smart-routing-for-2-cases-codex branch 2 times, most recently from c5e15dc to c8ec0f0 Compare September 4, 2026 23:47

lilly-luo marked this pull request as ready for review September 4, 2026 23:54

lilly-luo force-pushed the stack/lilly/only-do-smart-routing-for-2-cases-codex branch from c8ec0f0 to 354c435 Compare September 5, 2026 00:04

lilly-luo changed the title ~~only route codex / claude (no prompt) and codex -- '<>' and claude -- '<>' to smart router~~ Route only explicit Codex and Claude invocations through smart routing Sep 5, 2026

lilly-luo commented Sep 5, 2026

View reviewed changes

lilly-luo force-pushed the stack/lilly/only-do-smart-routing-for-2-cases-codex branch 3 times, most recently from 245a4a5 to 28bcd18 Compare September 5, 2026 00:17

lint

9cd4c7e

lilly-luo force-pushed the stack/lilly/only-do-smart-routing-for-2-cases-codex branch from 28bcd18 to 9cd4c7e Compare September 5, 2026 00:20

small tweak

678c43a

lilly-luo commented Sep 5, 2026

View reviewed changes

lilly-luo added 2 commits September 5, 2026 01:02

update

248df95

ruff

2dd0c67

lilly-luo mentioned this pull request Sep 5, 2026

fix nested config items in codex app #493

Closed

lilly-luo requested review from andy-xu-db, david-siqi-liu and rohita5l and removed request for david-siqi-liu and rohita5l September 5, 2026 01:52

lilly-luo requested a review from rohita5l September 5, 2026 01:52

Navigation Menu

Uh oh!

Route only explicit Codex and Claude invocations through smart routing - #502

Route only explicit Codex and Claude invocations through smart routing#502
lilly-luo wants to merge 6 commits into
mainfrom
stack/lilly/only-do-smart-routing-for-2-cases-codex

lilly-luo commented Sep 4, 2026 •

edited

Loading

Uh oh!

lilly-luo Sep 5, 2026

Uh oh!

lilly-luo Sep 5, 2026

Uh oh!

lilly-luo Sep 5, 2026

Uh oh!

lilly-luo Sep 5, 2026 •

edited

Loading

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

lilly-luo commented Sep 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

explanation

test

Uh oh!

lilly-luo Sep 5, 2026

Choose a reason for hiding this comment

Uh oh!

lilly-luo Sep 5, 2026

Choose a reason for hiding this comment

Uh oh!

lilly-luo Sep 5, 2026

Choose a reason for hiding this comment

Uh oh!

lilly-luo Sep 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

lilly-luo commented Sep 4, 2026 •

edited

Loading

lilly-luo Sep 5, 2026 •

edited

Loading