← Back to task

Commit 66af41df

commit 66af41df65f256e9ef14ac1fa1efad1295ca9afd
Author: Ben Sima <ben@bensima.com>
Date:   Thu Jan 1 17:34:38 2026

    Add Claude Code OAuth support for Claude Max/Pro (t-319)
    
    New Omni/Agent/Auth.hs module implements:
    - PKCE challenge generation (using sha256sum to avoid dep issues)
    - OAuth device code flow (print URL, user pastes code)
    - Token storage in ~/.local/share/agent/auth.json
    - Automatic token refresh when expired
    
    Usage:
      agent --provider=claude-code "hello"
    
    On first use, prompts user to visit auth URL and paste code.
    Subsequent uses automatically refresh the stored token.
    
    This allows using Claude without API costs for subscribers.
    
    Task-Id: t-319

diff --git a/Omni/Agent.hs b/Omni/Agent.hs
index a7eb1bf2..228858ac 100644
--- a/Omni/Agent.hs
+++ b/Omni/Agent.hs
@@ -33,6 +33,7 @@ import qualified Data.Aeson.Encode.Pretty as AesonPretty
 import qualified Data.ByteString.Lazy as BL
 import qualified Data.Text as Text
 import qualified Data.Text.IO as TextIO
+import qualified Omni.Agent.Auth as Auth
 import qualified Omni.Agent.Engine as Engine
 import qualified Omni.Agent.Provider as Provider
 import qualified Omni.Agent.Tools as Tools
@@ -290,6 +291,7 @@ resolveProvider providerName mModel = do
     "auto" -> resolveAutoProvider mModel
     "anthropic" -> resolveAnthropicProvider mModel
     "openrouter" -> resolveOpenRouterProvider mModel
+    "claude-code" -> resolveClaudeCodeProvider mModel
     "ollama" -> pure <| Right <| Provider.ollamaProvider (fromMaybe "llama3" mModel)
     other -> pure <| Left <| "Unknown provider: " <> other
 
@@ -327,3 +329,22 @@ resolveOpenRouterProvider mModel = do
   case mKey of
     Nothing -> pure <| Left "OPENROUTER_API_KEY not set"
     Just key -> pure <| Right <| Provider.openRouterProvider (Text.pack key) (fromMaybe "anthropic/claude-sonnet-4" mModel)
+
+-- | Resolve Claude Code provider (OAuth for Claude Max/Pro)
+resolveClaudeCodeProvider :: Maybe Text -> IO (Either Text Provider.Provider)
+resolveClaudeCodeProvider mModel = do
+  tokenResult <- Auth.getValidToken
+  case tokenResult of
+    Right token -> do
+      TextIO.hPutStrLn IO.stderr "[agent] Using Claude Code (OAuth)"
+      pure <| Right <| Provider.anthropicProvider token (fromMaybe "claude-sonnet-4-20250514" mModel)
+    Left err -> do
+      -- Need to login
+      TextIO.hPutStrLn IO.stderr <| "[agent] " <> err
+      TextIO.hPutStrLn IO.stderr "[agent] Starting OAuth login flow..."
+      loginResult <- Auth.loginAnthropic
+      case loginResult of
+        Left loginErr -> pure <| Left <| "OAuth login failed: " <> loginErr
+        Right creds -> do
+          TextIO.hPutStrLn IO.stderr "[agent] Login successful!"
+          pure <| Right <| Provider.anthropicProvider (Auth.credAccessToken creds) (fromMaybe "claude-sonnet-4-20250514" mModel)
diff --git a/Omni/Agent/Auth.hs b/Omni/Agent/Auth.hs
new file mode 100644
index 00000000..1367277e
--- /dev/null
+++ b/Omni/Agent/Auth.hs
@@ -0,0 +1,338 @@
+{-# LANGUAGE DeriveGeneric #-}
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE NoImplicitPrelude #-}
+
+-- | OAuth authentication for Claude Max/Pro subscriptions.
+--
+-- Implements the device code flow:
+-- 1. Generate PKCE challenge
+-- 2. Print auth URL for user to visit
+-- 3. User authorizes and gets code#oauthState
+-- 4. User pastes code back
+-- 5. Exchange for tokens
+-- 6. Store tokens in ~/.local/share/agent/auth.json
+--
+-- : out omni-agent-auth
+-- : dep aeson
+-- : dep http-conduit
+-- : dep directory
+-- : dep random
+-- : dep process
+module Omni.Agent.Auth
+  ( OAuthCredentials (..),
+    loginAnthropic,
+    refreshAnthropicToken,
+    loadCredentials,
+    saveCredentials,
+    getValidToken,
+    credentialsFile,
+    main,
+    test,
+  )
+where
+
+import Alpha
+import qualified Data.Aeson as Aeson
+import Data.Aeson ((.=), (.:))
+import qualified Data.Aeson.Key as AesonKey
+import qualified Data.Aeson.KeyMap as KeyMap
+import qualified Data.ByteString as BS
+
+import qualified System.Process as Process
+import qualified Data.ByteString.Lazy as BL
+import qualified Data.Text as Text
+import qualified Data.Text.Encoding as TE
+import qualified Data.Text.IO as TextIO
+import qualified Data.Time.Clock.POSIX as Time
+import qualified Network.HTTP.Simple as HTTP
+import qualified Omni.Test as Test
+import qualified System.Directory as Directory
+import qualified System.FilePath as FilePath
+import qualified System.IO as IO
+import System.Random (randomRIO)
+
+main :: IO ()
+main = Test.run test
+
+test :: Test.Tree
+test =
+  Test.group
+    "Omni.Agent.Auth"
+    [ Test.unit "PKCE generation produces valid output" <| do
+        (verifier, challenge) <- generatePKCE
+        -- Verifier should be 43 chars (32 bytes base64url encoded)
+        Text.length verifier Test.@=? 43
+        -- Challenge should also be 43 chars (SHA256 = 32 bytes)
+        Text.length challenge Test.@=? 43
+        -- They should be different
+        (verifier /= challenge) Test.@=? True,
+      Test.unit "credentialsFile returns valid path" <| do
+        path <- credentialsFile
+        -- Should end with auth.json
+        FilePath.takeFileName path Test.@=? "auth.json"
+    ]
+
+-- | OAuth credentials stored on disk
+data OAuthCredentials = OAuthCredentials
+  { credRefreshToken :: Text,
+    credAccessToken :: Text,
+    credExpiresAt :: Integer  -- Unix timestamp in milliseconds
+  }
+  deriving (Show, Eq, Generic)
+
+instance Aeson.ToJSON OAuthCredentials where
+  toJSON c =
+    Aeson.object
+      [ "refresh" .= credRefreshToken c,
+        "access" .= credAccessToken c,
+        "expires" .= credExpiresAt c
+      ]
+
+instance Aeson.FromJSON OAuthCredentials where
+  parseJSON = Aeson.withObject "OAuthCredentials" <| \v ->
+    OAuthCredentials
+      <$> v .: "refresh"
+      <*> v .: "access"
+      <*> v .: "expires"
+
+-- | Constants (same as pi)
+clientId :: Text
+clientId = "9d1c250a-e61b-44d9-88ed-5944d1962f5e"
+
+authorizeUrl :: Text
+authorizeUrl = "https://claude.ai/oauth/authorize"
+
+tokenUrl :: Text
+tokenUrl = "https://console.anthropic.com/v1/oauth/token"
+
+redirectUri :: Text
+redirectUri = "https://console.anthropic.com/oauth/code/callback"
+
+scopes :: Text
+scopes = "org:create_api_key user:profile user:inference"
+
+-- | Path to credentials file
+credentialsFile :: IO FilePath
+credentialsFile = do
+  dataDir <- Directory.getXdgDirectory Directory.XdgData "agent"
+  pure <| dataDir FilePath.</> "auth.json"
+
+-- | Generate PKCE verifier and challenge
+generatePKCE :: IO (Text, Text)
+generatePKCE = do
+  -- Generate 32 random bytes for verifier
+  verifierBytes <- replicateM 32 (randomRIO (0, 255) :: IO Int)
+  let verifier = base64UrlEncode <| BS.pack <| map fromIntegral verifierBytes
+  -- SHA256 hash of verifier for challenge (using sha256sum to avoid crypton dep issues)
+  hashHex <- Process.readProcess "sh" ["-c", "echo -n '" <> Text.unpack verifier <> "' | sha256sum | cut -d' ' -f1"] ""
+  let hashBytes = hexToBytes <| Text.pack <| filter (/= '\n') hashHex
+      challenge = base64UrlEncode hashBytes
+  pure (verifier, challenge)
+
+-- | Convert hex string to bytes
+hexToBytes :: Text -> BS.ByteString
+hexToBytes hex = BS.pack <| go (Text.unpack hex)
+  where
+    go [] = []
+    go [_] = []  -- odd length, ignore last
+    go (a:b:rest) = hexPair a b : go rest
+    hexPair a b = fromIntegral (hexDigit a * 16 + hexDigit b)
+    hexDigit c
+      | c >= '0' && c <= '9' = fromEnum c - fromEnum '0'
+      | c >= 'a' && c <= 'f' = fromEnum c - fromEnum 'a' + 10
+      | c >= 'A' && c <= 'F' = fromEnum c - fromEnum 'A' + 10
+      | otherwise = 0
+
+-- | Base64URL encode without padding (manual implementation to avoid dep issues)
+base64UrlEncode :: BS.ByteString -> Text
+base64UrlEncode bs = Text.pack <| go (BS.unpack bs)
+  where
+    -- Base64URL alphabet (- and _ instead of + and /)
+    alphabet :: Text
+    alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_"
+    
+    idx :: Int -> Char
+    idx i = Text.index alphabet i
+    
+    go :: [Word8] -> [Char]
+    go [] = []
+    go [a] =
+      let i0 = fromIntegral (a `shiftR` 2)
+          i1 = fromIntegral ((a .&. 0x03) `shiftL` 4)
+      in [idx i0, idx i1]
+    go [a, b] =
+      let i0 = fromIntegral (a `shiftR` 2)
+          i1 = fromIntegral (((a .&. 0x03) `shiftL` 4) .|. (b `shiftR` 4))
+          i2 = fromIntegral ((b .&. 0x0F) `shiftL` 2)
+      in [idx i0, idx i1, idx i2]
+    go (a : b : c : rest) =
+      let i0 = fromIntegral (a `shiftR` 2)
+          i1 = fromIntegral (((a .&. 0x03) `shiftL` 4) .|. (b `shiftR` 4))
+          i2 = fromIntegral (((b .&. 0x0F) `shiftL` 2) .|. (c `shiftR` 6))
+          i3 = fromIntegral (c .&. 0x3F)
+      in [idx i0, idx i1, idx i2, idx i3] ++ go rest
+
+-- | Login with Anthropic OAuth
+-- Returns credentials on success
+loginAnthropic :: IO (Either Text OAuthCredentials)
+loginAnthropic = do
+  (verifier, challenge) <- generatePKCE
+  
+  -- Build auth URL
+  let params =
+        [ ("code", "true"),
+          ("client_id", clientId),
+          ("response_type", "code"),
+          ("redirect_uri", redirectUri),
+          ("scope", scopes),
+          ("code_challenge", challenge),
+          ("code_challenge_method", "S256"),
+          ("oauthState", verifier)
+        ]
+      paramStr = Text.intercalate "&" <| map (\(k, v) -> k <> "=" <> v) params
+      authUrl = authorizeUrl <> "?" <> paramStr
+  
+  -- Print URL for user
+  TextIO.putStrLn ""
+  TextIO.putStrLn "To authenticate with Claude Max/Pro, visit this URL:"
+  TextIO.putStrLn ""
+  TextIO.putStrLn authUrl
+  TextIO.putStrLn ""
+  TextIO.putStrLn "After authorizing, you'll see a code. Paste it here:"
+  IO.hFlush IO.stdout
+  
+  -- Read code from user
+  authCode <- TextIO.getLine
+  let parts = Text.splitOn "#" authCode
+  case parts of
+    [code, oauthState] -> exchangeCode code oauthState verifier
+    [code] -> exchangeCode code verifier verifier  -- oauthState might be omitted
+    _ -> pure <| Left "Invalid code format. Expected: code#oauthState"
+
+-- | Exchange authorization code for tokens
+exchangeCode :: Text -> Text -> Text -> IO (Either Text OAuthCredentials)
+exchangeCode code oauthState verifier = do
+  let body =
+        Aeson.object
+          [ "grant_type" .= ("authorization_code" :: Text),
+            "client_id" .= clientId,
+            "code" .= code,
+            "oauthState" .= oauthState,
+            "redirect_uri" .= redirectUri,
+            "code_verifier" .= verifier
+          ]
+  
+  req0 <- HTTP.parseRequest (Text.unpack tokenUrl)
+  let req =
+        HTTP.setRequestMethod "POST"
+          <| HTTP.setRequestHeader "Content-Type" ["application/json"]
+          <| HTTP.setRequestBodyLBS (Aeson.encode body)
+          <| req0
+  
+  response <- HTTP.httpLBS req
+  let status = HTTP.getResponseStatusCode response
+      respBody = HTTP.getResponseBody response
+  
+  if status >= 200 && status < 300
+    then case Aeson.decode respBody of
+      Just obj -> do
+        let accessToken = fromMaybe "" <| lookupText "access_token" obj
+            refreshToken = fromMaybe "" <| lookupText "refresh_token" obj
+            expiresIn = fromMaybe 3600 <| lookupInt "expires_in" obj
+        -- Calculate expiry (current time + expires_in - 5 min buffer)
+        now <- round <$> Time.getPOSIXTime
+        let expiresAt = (now + expiresIn - 300) * 1000  -- Convert to milliseconds
+            creds = OAuthCredentials refreshToken accessToken expiresAt
+        -- Save credentials
+        saveCredentials creds
+        pure <| Right creds
+      Nothing -> pure <| Left <| "Failed to parse token response: " <> TE.decodeUtf8 (BL.toStrict respBody)
+    else pure <| Left <| "Token exchange failed: " <> TE.decodeUtf8 (BL.toStrict respBody)
+
+-- | Refresh an expired token
+refreshAnthropicToken :: Text -> IO (Either Text OAuthCredentials)
+refreshAnthropicToken refreshToken = do
+  let body =
+        Aeson.object
+          [ "grant_type" .= ("refresh_token" :: Text),
+            "client_id" .= clientId,
+            "refresh_token" .= refreshToken
+          ]
+  
+  req0 <- HTTP.parseRequest (Text.unpack tokenUrl)
+  let req =
+        HTTP.setRequestMethod "POST"
+          <| HTTP.setRequestHeader "Content-Type" ["application/json"]
+          <| HTTP.setRequestBodyLBS (Aeson.encode body)
+          <| req0
+  
+  response <- HTTP.httpLBS req
+  let status = HTTP.getResponseStatusCode response
+      respBody = HTTP.getResponseBody response
+  
+  if status >= 200 && status < 300
+    then case Aeson.decode respBody of
+      Just obj -> do
+        let accessToken = fromMaybe "" <| lookupText "access_token" obj
+            newRefreshToken = fromMaybe refreshToken <| lookupText "refresh_token" obj
+            expiresIn = fromMaybe 3600 <| lookupInt "expires_in" obj
+        now <- round <$> Time.getPOSIXTime
+        let expiresAt = (now + expiresIn - 300) * 1000
+            creds = OAuthCredentials newRefreshToken accessToken expiresAt
+        saveCredentials creds
+        pure <| Right creds
+      Nothing -> pure <| Left "Failed to parse refresh response"
+    else pure <| Left <| "Token refresh failed: " <> TE.decodeUtf8 (BL.toStrict respBody)
+
+-- | Load credentials from disk
+loadCredentials :: IO (Maybe OAuthCredentials)
+loadCredentials = do
+  path <- credentialsFile
+  exists <- Directory.doesFileExist path
+  if exists
+    then do
+      content <- BL.readFile path
+      pure <| Aeson.decode content
+    else pure Nothing
+
+-- | Save credentials to disk
+saveCredentials :: OAuthCredentials -> IO ()
+saveCredentials creds = do
+  path <- credentialsFile
+  Directory.createDirectoryIfMissing True (FilePath.takeDirectory path)
+  BL.writeFile path (Aeson.encode creds)
+
+-- | Get a valid token, refreshing if necessary
+-- Returns Left if login is required
+getValidToken :: IO (Either Text Text)
+getValidToken = do
+  mCreds <- loadCredentials
+  case mCreds of
+    Nothing -> pure <| Left "Not logged in. Run: agent --provider=claude-code login"
+    Just creds -> do
+      now <- round <$> Time.getPOSIXTime
+      let nowMs = now * 1000
+      if credExpiresAt creds > nowMs
+        then pure <| Right <| credAccessToken creds
+        else do
+          -- Token expired, try to refresh
+          result <- refreshAnthropicToken (credRefreshToken creds)
+          case result of
+            Left err -> pure <| Left <| "Token refresh failed: " <> err
+            Right newCreds -> pure <| Right <| credAccessToken newCreds
+
+-- | Helper to lookup text in JSON object
+lookupText :: Text -> Aeson.Value -> Maybe Text
+lookupText key (Aeson.Object obj) =
+  case KeyMap.lookup (AesonKey.fromText key) obj of
+    Just (Aeson.String s) -> Just s
+    _ -> Nothing
+lookupText _ _ = Nothing
+
+-- | Helper to lookup int in JSON object
+lookupInt :: Text -> Aeson.Value -> Maybe Integer
+lookupInt key (Aeson.Object obj) =
+  case KeyMap.lookup (AesonKey.fromText key) obj of
+    Just (Aeson.Number n) -> Just <| round n
+    _ -> Nothing
+lookupInt _ _ = Nothing
diff --git a/Omni/Bild/Deps/Haskell.nix b/Omni/Bild/Deps/Haskell.nix
index e4830bb8..f562d167 100644
--- a/Omni/Bild/Deps/Haskell.nix
+++ b/Omni/Bild/Deps/Haskell.nix
@@ -8,6 +8,7 @@
   "aeson"
   "async"
   "base"
+  "base64-bytestring"
   "bytestring"
   "clay"
   "cmark"
@@ -16,6 +17,7 @@
   "conduit-extra"
   "config-ini"
   "containers"
+  "crypton"
   "directory"
   "docopt"
   "envy"