Rewrite Agent Loop as Native Op

Rewrite the core agent loop using native Op primitives instead of Engine.hs.

Context

Once Op.hs, Trace.hs, and Sequential interpreter are working, we can express the agent loop directly as an Op program. This replaces the monolithic Engine.runAgentWithProvider.

Read Omni/Agent/ARCHITECTURE.md for full design rationale.

Deliverables

1. Create Omni/Agent/Programs/Agent.hs

module Omni.Agent.Programs.Agent where

-- | Standard agent loop as an Op
-- This replaces Engine.runAgentWithProvider
agent :: AgentConfig -> Text -> Op AgentState AgentResult
agent config userPrompt = do
  emit (EventCustom "agent_start" (toJSON config))
  checkpoint "init"
  
  -- Initialize state
  put initialState
  
  -- Build initial prompt
  let systemMsg = Message System (agentSystemPrompt config)
      userMsg = Message User userPrompt
      initialPrompt = Prompt [systemMsg, userMsg]
  
  -- Run loop
  result <- limit (agentBudget config) (agentLoop config initialPrompt 0)
  
  case result of
    Left exhausted -> do
      emit (EventExhausted (tshow exhausted) <timestamp>)
      s <- get
      pure AgentResult
        { resultFinalMessage = ""
        , resultError = Just ("Budget exhausted: " <> tshow exhausted)
        , resultIterations = asIteration s
        , resultTotalCost = asTotalCost s
        , resultTotalTokens = asTotalTokens s
        , resultToolCallCount = asToolCalls s
        }
    Right response -> do
      emit (EventCustom "agent_complete" (toJSON response))
      s <- get
      pure AgentResult
        { resultFinalMessage = response
        , resultError = Nothing
        , resultIterations = asIteration s
        , resultTotalCost = asTotalCost s
        , resultTotalTokens = asTotalTokens s
        , resultToolCallCount = asToolCalls s
        }

-- | The core agent loop
agentLoop :: AgentConfig -> Prompt -> Int -> Op AgentState Text
agentLoop config prompt iteration = do
  -- Check steering
  steering <- check
  case steering of
    Just SteeringCancel -> pure "Cancelled by user"
    Just (SteeringMessage msg) -> do
      -- Inject steering message and continue
      let prompt' = appendMessage prompt (Message User msg)
      agentLoop config prompt' iteration
    Nothing -> do
      -- Update iteration
      modify (\s -> s { asIteration = iteration })
      
      -- Checkpoint before expensive inference
      when (iteration > 0 && iteration `mod` 5 == 0) $
        checkpoint ("iteration-" <> tshow iteration)
      
      -- Call LLM
      response <- infer (configModel config) prompt
      
      -- Update cost tracking
      modify (\s -> s 
        { asTotalCost = asTotalCost s + responseCost response
        , asTotalTokens = asTotalTokens s + responseTokens response
        })
      
      -- Check for tool calls
      case responseToolCalls response of
        [] -> do
          -- Done - return final message
          pure (responseContent response)
        
        toolCalls -> do
          -- Execute tools
          results <- executeTools (configTools config) toolCalls
          
          -- Update tool call count
          modify (\s -> s { asToolCalls = asToolCalls s + length toolCalls })
          
          -- Build new prompt with tool results
          let prompt' = prompt 
                 assistantMessage response
                 toolResultMessages results
          
          -- Check iteration limit
          if iteration >= configMaxIterations config
            then pure (responseContent response)
            else agentLoop config prompt' (iteration + 1)

-- | Execute tool calls
executeTools :: [Tool] -> [ToolCall] -> Op AgentState [ToolResult]
executeTools tools calls = traverse (executeTool tools) calls

executeTool :: [Tool] -> ToolCall -> Op AgentState ToolResult
executeTool tools tc = do
  let name = tcName tc
      args = tcArgs tc
  case find (\t -> toolName t == name) tools of
    Nothing -> pure (ToolResult (tcId tc) ("Tool not found: " <> name) False)
    Just t -> do
      result <- tool name args
      pure (ToolResult (tcId tc) (renderResult result) True)

2. Agent State

data AgentState = AgentState
  { asHistory :: [Message]      -- conversation history
  , asTotalCost :: Double       -- cents
  , asTotalTokens :: Int
  , asIteration :: Int
  , asToolCalls :: Int
  }
  deriving (Show, Eq, Generic)

instance ToJSON AgentState
instance FromJSON AgentState

-- For CRDT support (parallel agents)
instance Semigroup AgentState where
  a <> b = AgentState
    { asHistory = asHistory a <> asHistory b  -- or: take longer?
    , asTotalCost = max (asTotalCost a) (asTotalCost b)  -- LWW-ish
    , asTotalTokens = max (asTotalTokens a) (asTotalTokens b)
    , asIteration = max (asIteration a) (asIteration b)
    , asToolCalls = asToolCalls a + asToolCalls b
    }

initialState :: AgentState
initialState = AgentState [] 0 0 0 0

3. Convenience Function

-- | Drop-in replacement for Engine.runAgentWithProvider
runAgent 
  :: Provider 
  -> AgentConfig 
  -> Text 
  -> IO (Either Text AgentResult)
runAgent provider config prompt = do
  let program = agent config prompt
      seqConfig = defaultSeqConfig provider (configTools config)
  result <- runSequential seqConfig initialState program
  case result of
    Left err -> pure (Left err)
    Right (a, _, _) -> pure (Right a)

4. Update Exports

In Omni/Agent.hs, add:

import qualified Omni.Agent.Programs.Agent as Programs

-- Optionally re-export the new version
runAgentOp :: Provider -> AgentConfig -> Text -> IO (Either Text AgentResult)
runAgentOp = Programs.runAgent

Notes

• This should produce identical behavior to Engine.runAgentWithProvider
• But now it's composable - can use par, race, etc.
• Checkpoints at iteration boundaries for long runs
• State is explicit and serializable
• Traces capture everything

Testing

• Compare output with Engine.runAgentWithProvider for same inputs
• Verify cost/token tracking matches
• Verify tool execution works
• Verify iteration limit enforced
• Verify checkpoints saved at correct intervals
• Test steering (cancel, inject message)

Files to Read First

• Omni/Agent/ARCHITECTURE.md
• Omni/Agent/Op.hs
• Omni/Agent/Interpreter/Sequential.hs
• Omni/Agent/Engine.hs (current implementation to match)