diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,32 @@ ## Unreleased +## 0.4.0.0 — 2026-09-08 + +- Raise the internal `shikumi` bound to `^>=0.4.0.0`, and give the test suite's `shikumi-cache` dependency an explicit `^>=0.2.0.0` bound where it previously had none. + +- Verify released Responses wire schemas, exact opaque checkpoint replay, call IDs, origin rejection, terminal errors, retries, cache observations and server cleanup with loopback integration tests. + +- Write version-2 reasoning-safe checkpoints while reading version 1. Add safe summary rendering, explicit-model startup and pure restart; defer unsafe compaction and reject incompatible continuation before dispatch. Public API additions require a PVP review at release. + +- Preserve structured refusals in defensive ReAct response checks and use shared error rendering for tool observations; partial refused tool calls never dispatch. + +- Upgrade the dependency on `mori://shinzui/baikai/packages/baikai` to `>=0.7.0.0 && <0.8`. Preserve cost basis and usage availability in agent checkpoints, including exact rational costs and legacy checkpoint decoding. + +### Added + +- Rich `ToolOutput` results and runtime-schema dynamic tool registration, with + compatibility text projections. +- Completed-turn ReAct sessions with original native call IDs, validated final + submission, versioned checkpoints, continuation, and whole-exchange compaction. + Legacy ReAct extraction and CodeAct trajectory APIs remain unchanged. + +### Changed + +- Widened the `http-client-tls` upper bound to admit 0.4. The web tool uses + the unchanged `tlsManagerSettings` API, and 0.4 is the compatible line for + the TLS 2.x / Crypton 1.1 cohort selected by baikai 0.6. + ## 0.3.0.3 — 2026-08-29 ### Changed diff --git a/shikumi-tools.cabal b/shikumi-tools.cabal --- a/shikumi-tools.cabal +++ b/shikumi-tools.cabal @@ -1,6 +1,6 @@ cabal-version: 3.4 name: shikumi-tools -version: 0.3.0.3 +version: 0.4.0.0 synopsis: Typed tools and ReAct agents for shikumi LM programs (EP-11) @@ -43,22 +43,26 @@ import: common-options hs-source-dirs: src exposed-modules: + Shikumi.Agent.History Shikumi.Agent.ReAct Shikumi.CodeExec.CodeAct Shikumi.CodeExec.Interpreter Shikumi.CodeExec.ProgramOfThought Shikumi.CodeExec.Prompt + Shikumi.CodeExec.RLM + Shikumi.CodeExec.Session Shikumi.Tool Shikumi.Tool.Builtin Shikumi.Tool.Builtin.Fs Shikumi.Tool.Builtin.Shell Shikumi.Tool.Builtin.Web Shikumi.Tool.Env + Shikumi.Tool.Output Shikumi.Tool.Web build-depends: , aeson >=2.2 && <2.3 - , baikai >=0.6 && <0.7 + , baikai >=0.7.0.0 && <0.8 , base >=4.20 && <5 , bytestring >=0.11 && <0.13 , containers >=0.6 && <0.9 @@ -67,12 +71,12 @@ , filepath >=1.4 && <1.6 , generic-lens >=2.2 && <2.4 , http-client >=0.7 && <0.8 - , http-client-tls >=0.3 && <0.4 + , http-client-tls >=0.3 && <0.5 , http-types >=0.12 && <0.13 , lens ^>=5.3 , process >=1.6 && <1.7 , regex-tdfa >=1.3 && <1.4 - , shikumi ^>=0.3.0.0 + , shikumi ^>=0.4.0.0 , text ^>=2.1 , vector >=0.13 && <0.14 @@ -84,25 +88,32 @@ ghc-options: -threaded -with-rtsopts=-N other-modules: AcceptanceSpec + AgentHistorySpec BuiltinAcceptanceSpec CodeActSpec CompactionSpec EnvSpec Fixtures FsSpec - MockLLM ProgramOfThoughtSpec ProtocolSpec + ReActSessionExample ReActSpec + ResponsesIntegrationSpec + ResponsesSchemaSpec RestrictedSpec + RLMExample + RLMSpec SchemaSpec + SessionSpec ShellSpec + ToolOutputSpec ToolSpec WebSpec build-depends: , aeson - , baikai >=0.6 && <0.7 + , baikai >=0.7.0.0 && <0.8 , base , bytestring , containers @@ -111,8 +122,10 @@ , filepath , generic-lens , lens - , shikumi ^>=0.3.0.0 - , shikumi-tools ^>=0.3.0.0 + , shikumi ^>=0.4.0.0 + , shikumi-cache ^>=0.2.0.0 + , shikumi-testing ^>=0.1.0.0 + , shikumi-tools ^>=0.4.0.0 , tasty , tasty-hunit , text diff --git a/src/Shikumi/Agent/History.hs b/src/Shikumi/Agent/History.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/Agent/History.hs @@ -0,0 +1,338 @@ +-- | Versioned, completed-exchange checkpoints. Registries and closures stay caller-owned. +module Shikumi.Agent.History + ( ReActSession, + HistoryError (..), + HistoryEntry (..), + newSession, + sessionOrigin, + sessionPrefix, + bindSessionOrigin, + protectSession, + compactionSafe, + renderSummaryEntry, + renderSessionSummaryInput, + restartSessionFromSummary, + sessionProtocol, + sessionFingerprint, + sessionTools, + sessionTurns, + sessionIterations, + sessionFinished, + sessionCompactedThrough, + auditHistory, + promptEntries, + promptMessages, + entryMessages, + appendUser, + appendExchange, + validateCalls, + parsePromptActions, + validateSession, + compactSession, + encodeSession, + decodeSession, + ) +where + +import Baikai qualified as B +import Baikai.Cost qualified as BC +import Control.Lens ((&), (.~), (^.)) +import Control.Monad (foldM, unless, when) +import Data.Aeson +import Data.Aeson.Types (Parser, parseEither) +import Data.Bifunctor (first) +import Data.ByteString.Lazy qualified as LBS +import Data.Generics.Labels () +import Data.Ratio (denominator, numerator, (%)) +import Data.Set qualified as Set +import Data.Text (Text) +import Data.Text qualified as T +import Data.Text.Encoding (decodeUtf8, encodeUtf8) +import Data.Vector qualified as V +import Shikumi.LLM.Continuation (RequestOrigin, contextIdentity, hasOpaqueContinuation) +import Shikumi.Tool.Output + +newtype HistoryError = HistoryError Text deriving stock (Eq, Show) + +-- | Rejected responses are audit-only; their correction replaces them in requests. +-- Accepted exchanges contain the original assistant payload and ordered results. +data HistoryEntry + = UserTurn !B.UserPayload + | Exchange !B.AssistantPayload ![(B.ToolCall, ToolOutput)] !(Maybe Text) + deriving stock (Eq, Show) + +data ReActSession = ReActSession + { sessionProtocol :: !Text, + sessionFingerprint :: !Value, + sessionTools :: !Value, + sessionTurns :: !Int, + sessionIterations :: !Int, + sessionFinished :: !Bool, + auditHistory :: ![HistoryEntry], + compactedThrough :: !Int, + summary :: !(Maybe Text), + sessionOrigin :: !(Maybe RequestOrigin), + sessionPrefix :: !(Maybe Value) + } + deriving stock (Eq, Show) + +newSession :: Text -> Value -> Value -> Text -> Either HistoryError ReActSession +newSession proto fingerprint tools input = do + let s = ReActSession proto fingerprint tools 0 0 False [userEntry input] 0 Nothing Nothing Nothing + validateSession s + pure s + +userEntry :: Text -> HistoryEntry +userEntry t = UserTurn (B.UserPayload (V.singleton (B.UserText (B.TextContent t))) Nothing) + +sessionCompactedThrough :: ReActSession -> Int +sessionCompactedThrough = compactedThrough + +appendUser :: Text -> ReActSession -> Either HistoryError ReActSession +appendUser t s = do + validateSession s + pure s {auditHistory = auditHistory s <> [userEntry t], sessionIterations = 0, sessionFinished = False} + +callsOf :: B.AssistantPayload -> [B.ToolCall] +callsOf p = [c | B.AssistantToolCall c <- V.toList (p ^. #content)] + +-- | Validate the whole native proposal before dispatch, including prior accepted IDs. +validateCalls :: [Text] -> [B.ToolCall] -> Either HistoryError () +validateCalls previous calls = do + _ <- foldM check (Set.fromList previous) calls + pure () + where + check seen c + | T.null (T.strip (c ^. #id_)) = Left (HistoryError "empty tool call ID") + | Set.member (c ^. #id_) seen = Left (HistoryError "duplicate tool call ID") + | B.isCutOffToolCall c = Left (HistoryError "cut-off tool arguments") + | otherwise = Right (Set.insert (c ^. #id_) seen) + +appendExchange :: B.AssistantPayload -> [(B.ToolCall, ToolOutput)] -> Maybe Text -> Bool -> ReActSession -> Either HistoryError ReActSession +appendExchange p outputs correction finished s = do + let next = s {auditHistory = auditHistory s <> [Exchange p outputs correction], sessionTurns = sessionTurns s + 1, sessionIterations = sessionIterations s + 1, sessionFinished = finished} + validateSession next + pure next + +validateSession :: ReActSession -> Either HistoryError () +validateSession s = do + ensure (sessionProtocol s `elem` ["native", "prompt"]) "unknown protocol" + ensure (not (null (auditHistory s))) "empty history" + case auditHistory s of + UserTurn _ : _ -> pure () + _ -> Left (HistoryError "history must start with a user turn") + ensure (sessionTurns s == length [() | Exchange {} <- auditHistory s]) "turn count mismatch" + ensure (sessionIterations s >= 0 && sessionIterations s <= sessionTurns s) "invalid iteration count" + ensure (compactedThrough s >= 0 && compactedThrough s <= length (auditHistory s)) "invalid compaction boundary" + ensure ((compactedThrough s == 0) == (summary s == Nothing)) "missing compaction summary" + ensure (sessionIterations s == length (takeWhile isExchange (reverse (auditHistory s)))) "iteration count mismatch" + when (sessionFinished s) $ case reverse (auditHistory s) of + Exchange _ [(call, _)] Nothing : _ -> ensure (call ^. #name == "shikumi_submit_final") "finished without final submission" + _ -> Left (HistoryError "finished without accepted exchange") + _ <- foldM validateEntry [] (auditHistory s) + pure () + where + isExchange Exchange {} = True + isExchange _ = False + validateEntry ids (UserTurn _) = pure ids + validateEntry ids (Exchange p outputs correction) = case correction of + Just reason -> do + ensure (not (T.null reason) && null outputs) "rejected exchange has results or no correction" + pure ids + Nothing -> do + let calls = map fst outputs + ensure (not (null calls)) "unresolved exchange" + validateCalls ids calls + ensure (not (any ((== "shikumi_submit_final") . (^. #name)) calls) || length calls == 1) "mixed final submission" + when (sessionProtocol s == "native") $ + ensure (callsOf p == calls) "results do not match assistant calls in order" + when (sessionProtocol s == "prompt") $ do + ensure (null (callsOf p)) "prompt exchange contains native calls" + ensure (all (T.isPrefixOf "prompt-" . (^. #id_)) calls) "invalid synthetic prompt ID" + let text = T.concat [t | B.AssistantText (B.TextContent t) <- V.toList (p ^. #content)] + actions <- parsePromptActions text + ensure (actions == map (\c -> (c ^. #name, c ^. #arguments)) calls) "prompt results do not match actions" + pure (ids <> map (^. #id_) calls) + +-- | Shared proposal/checkpoint parser so acceptance cannot diverge after dispatch. +parsePromptActions :: Text -> Either HistoryError [(Text, Value)] +parsePromptActions text = first (HistoryError . T.pack) $ do + let trimmed = T.strip text + stripped = case T.stripPrefix "```" trimmed of + Nothing -> trimmed + Just rest -> T.strip (fst (T.breakOn "```" (T.drop 1 (T.dropWhile (/= '\n') rest)))) + value <- eitherDecodeStrict (encodeUtf8 stripped) + parseEither (withObject "proposal" (\o -> o .: "calls" >>= traverse (withObject "call" (\c -> (,) <$> c .: "tool" <*> c .: "args")))) value + +ensure :: Bool -> Text -> Either HistoryError () +ensure ok message = unless ok (Left (HistoryError message)) + +entryMessages :: Text -> HistoryEntry -> [B.Message] +entryMessages _ (UserTurn p) = [B.UserMessage p] +entryMessages _ (Exchange _ _ (Just correction)) = [B.user ("Protocol correction: " <> correction)] +entryMessages proto (Exchange p outputs Nothing) + | proto == "native" = B.AssistantMessage p : map (uncurry toolOutputMessage) outputs + | otherwise = B.AssistantMessage p : [B.user (T.intercalate "\n" [c ^. #name <> ": " <> renderToolOutput o | (c, o) <- outputs])] + +promptEntries :: ReActSession -> [HistoryEntry] +promptEntries s = maybe [] (\t -> [userEntry ("Earlier conversation summary: " <> t)]) (summary s) <> drop (compactedThrough s) (auditHistory s) + +promptMessages :: ReActSession -> [B.Message] +promptMessages s = concatMap (entryMessages (sessionProtocol s)) (promptEntries s) + +-- | Replace only an old prefix of complete entries; the full audit is unchanged. +compactSession :: Int -> Text -> ReActSession -> Either HistoryError ReActSession +compactSession through text s = do + ensure (compactionSafe through s) "compaction would change the prefix of opaque reasoning; explicitly restart from a caller-approved summary" + ensure (through > compactedThrough s) "compaction must advance" + let next = s {compactedThrough = through, summary = Just text} + validateSession next + pure next + +-- | Any retained opaque exchange still depends on the original prefix. +compactionSafe :: Int -> ReActSession -> Bool +compactionSafe through s = not (hasOpaqueContinuation (concatMap (entryMessages (sessionProtocol s)) (drop through (auditHistory s)))) + +bindSessionOrigin :: Maybe RequestOrigin -> ReActSession -> ReActSession +bindSessionOrigin origin s = s {sessionOrigin = origin} + +-- | Capture the exact request view after completing the exchange. This is +-- compatibility data, not an authenticity proof for untrusted checkpoints. +protectSession :: B.Context -> ReActSession -> ReActSession +protectSession ctx s = s {sessionPrefix = if hasOpaqueContinuation (promptMessages s) then Just (contextIdentity ctx) else Nothing} + +-- | Structural projection: provider-owned reasoning and unknown rich extensions +-- never enter a summarizer request. The full audit retains all original bytes. +renderSummaryEntry :: HistoryEntry -> Text +renderSummaryEntry (UserTurn p) = T.intercalate "\n" [case b of B.UserText (B.TextContent t) -> t; B.UserImage _ -> "[image omitted]" | b <- V.toList (p ^. #content)] +renderSummaryEntry (Exchange p outputs correction) = T.intercalate "\n" (map block (V.toList (p ^. #content)) <> map output outputs <> maybe [] (\t -> ["Correction: " <> t]) correction) + where + block (B.AssistantText (B.TextContent t)) = t + block (B.AssistantThinking t) + | t ^. #redacted = "[redacted reasoning omitted]" + | otherwise = t ^. #thinking <> "\n[reasoning continuation omitted]" + block (B.AssistantToolCall c) = c ^. #name <> " " <> jsonText (c ^. #arguments) + output (c, o) = + c ^. #name + <> ": " + <> T.intercalate + "\n" + ( (if result o ^. #isError then ["Tool error"] else []) + <> [case b of B.ToolResultText (B.TextContent t) -> t; B.ToolResultImage _ -> "[image omitted]" | b <- V.toList (result o ^. #content)] + <> maybe [] (\v -> ["Structured JSON: " <> jsonText v]) (structuredContent o) + <> ["[extension blocks omitted]" | not (null (extensionBlocks o))] + ) + jsonText = decodeUtf8 . LBS.toStrict . encode + +renderSessionSummaryInput :: ReActSession -> Text +renderSessionSummaryInput = T.intercalate "\n" . map renderSummaryEntry . auditHistory + +-- | A separate, unbound conversation; no IO and no old native exchanges. +restartSessionFromSummary :: Text -> ReActSession -> Either HistoryError ReActSession +restartSessionFromSummary text s = do + validateSession s + ensure (not (T.null (T.strip text))) "restart summary must not be empty" + newSession (sessionProtocol s) (sessionFingerprint s) (sessionTools s) ("Earlier conversation summary: " <> text) + +-- The transfer format is local and explicit, independent of Baikai's Message JSON. +encodeSession :: ReActSession -> Value +encodeSession s = + object + [ "version" .= (2 :: Int), + "origin" .= sessionOrigin s, + "protectedPrefix" .= sessionPrefix s, + "protocol" .= sessionProtocol s, + "fingerprint" .= sessionFingerprint s, + "tools" .= sessionTools s, + "turns" .= sessionTurns s, + "iterations" .= sessionIterations s, + "finished" .= sessionFinished s, + "history" .= map encodeEntry (auditHistory s), + "compactedThrough" .= compactedThrough s, + "summary" .= summary s + ] + +decodeSession :: Value -> Either HistoryError ReActSession +decodeSession value = do + s <- first (HistoryError . T.pack) (parseEither parseSession value) + validateSession s + pure s + +parseSession :: Value -> Parser ReActSession +parseSession = withObject "ReActSession" $ \o -> do + version <- o .: "version" + unless (version `elem` [1, 2 :: Int]) (fail "unknown checkpoint version") + ReActSession + <$> o .: "protocol" + <*> o .: "fingerprint" + <*> o .: "tools" + <*> o .: "turns" + <*> o .: "iterations" + <*> o .: "finished" + <*> (o .: "history" >>= traverse parseEntry) + <*> o .: "compactedThrough" + <*> o .: "summary" + <*> (if version == 1 then pure Nothing else o .: "origin") + <*> (if version == 1 then pure Nothing else o .: "protectedPrefix") + +encodeEntry :: HistoryEntry -> Value +encodeEntry (UserTurn p) = object ["kind" .= String "user", "content" .= (p ^. #content), "timestamp" .= (p ^. #timestamp)] +encodeEntry (Exchange p outputs correction) = + object + [ "kind" .= String "exchange", + "content" .= (p ^. #content), + "timestamp" .= (p ^. #timestamp), + "stop" .= (p ^. #stopReason), + "error" .= (p ^. #errorMessage), + "usage" .= encodeUsage (p ^. #usage), + "results" .= [object ["call" .= c, "blocks" .= (result r ^. #content), "isError" .= (result r ^. #isError), "structured" .= structuredContent r, "extensions" .= extensionBlocks r] | (c, r) <- outputs], + "correction" .= correction + ] + +parseEntry :: Value -> Parser HistoryEntry +parseEntry = withObject "HistoryEntry" $ \o -> do + kind <- o .: "kind" :: Parser Text + case kind of + "user" -> UserTurn <$> (B.UserPayload <$> o .: "content" <*> o .: "timestamp") + "exchange" -> + Exchange + <$> (B.AssistantPayload <$> o .: "content" <*> (o .: "usage" >>= parseUsage) <*> o .: "stop" <*> o .: "error" <*> o .: "timestamp") + <*> ( o .: "results" + >>= traverse + ( withObject "result" $ \r -> + (,) <$> r .: "call" <*> (ToolOutput <$> (B.ToolResult <$> r .: "blocks" <*> r .: "isError") <*> r .: "structured" <*> r .: "extensions") + ) + ) + <*> o .: "correction" + _ -> fail "unknown history entry" + +-- Costs use integer ratios, preserving even non-terminating rational values exactly. +encodeUsage :: B.Usage -> Value +encodeUsage u = + object + [ "tokens" .= toJSON u, + "costRatios" + .= map + (\r -> (numerator r, denominator r)) + [c ^. #usd, b ^. #inputUsd, b ^. #outputUsd, b ^. #cachedInputUsd, b ^. #cachedWriteUsd] + ] + where + c = u ^. #cost + b = c ^. #breakdown + +parseUsage :: Value -> Parser B.Usage +parseUsage = withObject "usage" $ \o -> do + ratios <- o .: "costRatios" :: Parser [(Integer, Integer)] + unless (all ((> 0) . snd) ratios) (fail "invalid cost denominator") + basis <- o .: "tokens" >>= withObject "tokens" (\t -> t .: "cost" >>= withObject "cost" (\c -> c .:? "basis" .!= mempty)) + cost <- case map (uncurry (%)) ratios of + [total, i, out, cached, written] -> pure (BC.Cost total (BC.CostBreakdown i out cached written) basis) + _ -> fail "invalid cost breakdown" + u <- + o .: "tokens" + >>= withObject + "tokens" + ( \t -> + B.Usage <$> t .: "input_tokens" <*> t .: "output_tokens" <*> t .: "cache_read_tokens" <*> t .: "cache_write_tokens" <*> t .: "reasoning_tokens" <*> t .: "total_tokens" <*> t .:? "availability" <*> pure cost + ) + pure (u & #cost .~ cost) diff --git a/src/Shikumi/Agent/ReAct.hs b/src/Shikumi/Agent/ReAct.hs --- a/src/Shikumi/Agent/ReAct.hs +++ b/src/Shikumi/Agent/ReAct.hs @@ -36,6 +36,13 @@ -- * Building agents react, reactWithTrajectory, + SessionResult (..), + finalToolName, + startSession, + startSessionWithModel, + advanceSession, + continueSession, + runSession, -- * The protocol seam Proposal (..), @@ -69,8 +76,9 @@ flattenAssistantBlocks, user, ) +import Baikai qualified as B import Control.Lens ((&), (.~), (^.)) -import Data.Aeson (Value (..), eitherDecodeStrict, encode) +import Data.Aeson (Value (..), eitherDecodeStrict, encode, object, (.=)) import Data.Aeson.KeyMap qualified as KM import Data.ByteString.Lazy qualified as LBS import Data.Generics.Labels () @@ -86,13 +94,17 @@ import Effectful.Error.Static (Error, catchError, throwError) import GHC.Generics (Generic) import Shikumi.Adapter (ModelCapability (..), ToPrompt (toPrompt), attachSchema, capabilityFor) +import Shikumi.Agent.History (ReActSession) +import Shikumi.Agent.History qualified as H import Shikumi.Compaction (CompactionConfig (..), compactTail, defaultCompactionConfig, usageExceedsWindow) -import Shikumi.Error (ShikumiError (..)) +import Shikumi.Error (ShikumiError (..), fromBaikaiError) import Shikumi.LLM (LLM, complete) +import Shikumi.LLM.Continuation qualified as C import Shikumi.Program (Program (FMap), embed) -import Shikumi.Schema (FromModel, ToSchema, Validatable, parseOutput, toSchema) +import Shikumi.Schema (FromModel, ToSchema, Validatable, fromModelChecked, parseOutput, toSchema) import Shikumi.Signature (Signature, getInstruction) -import Shikumi.Tool (ToolRegistry, registryBaikai, registryTools, renderToolError, runToolCall, someToolDescription, someToolName, someToolSchema) +import Shikumi.Tool (ToolRegistry, registryBaikai, registryNames, registryTools, renderToolError, runToolCall, runToolCallOutput, someToolDescription, someToolName, someToolSchema) +import Shikumi.Tool.Output (ToolOutput (..), textToolOutput) -- --------------------------------------------------------------------------- -- The trajectory data model @@ -574,3 +586,189 @@ -- | Compact-encode a JSON value to text. encodeText :: Value -> Text encodeText = decodeUtf8 . LBS.toStrict . encode + +-- --------------------------------------------------------------------------- +-- Completed-turn sessions (the legacy loop above retains its extraction call) +-- --------------------------------------------------------------------------- + +-- | A pause always contains complete exchanges, including all dispatched results. +data SessionResult o = SessionPaused ReActSession | SessionFinished o ReActSession + deriving stock (Eq, Show) + +finalToolName :: Text +finalToolName = "shikumi_submit_final" + +-- | The fingerprint stores exact schema and instruction values (no lossy hash). +sessionIdentity :: forall i o. (ToSchema o) => Signature i o -> Value +sessionIdentity sig = + object + [ "instruction" .= getInstruction sig, + "inputFields" .= show (sig ^. #inputFields), + "outputSchema" .= toSchema (Proxy @o) + ] + +sessionToolSchemas :: ToolRegistry -> Value +sessionToolSchemas reg = + object + [ "tools" .= [object ["name" .= someToolName t, "schema" .= someToolSchema t, "description" .= someToolDescription t] | t <- registryTools reg] + ] + +historyOrThrow :: (Error ShikumiError :> es) => Either H.HistoryError a -> Eff es a +historyOrThrow = either (\(H.HistoryError t) -> throwError (ValidationFailure ("ReAct history: " <> t))) pure + +protocolText :: ReActConfig -> Text +protocolText cfg = case resolveProtocolKind (protocol cfg) emptyModel of + ProtocolNative -> "native" + _ -> "prompt" + +checkSession :: (ToSchema o, Error ShikumiError :> es) => Signature i o -> ToolRegistry -> ReActConfig -> ReActSession -> Eff es () +checkSession sig reg cfg s = do + historyOrThrow (H.validateSession s) + if finalToolName `elem` registryNames reg + then throwError (ValidationFailure "Registry collides with shikumi_submit_final") + else pure () + if H.sessionFingerprint s == sessionIdentity sig && H.sessionTools s == sessionToolSchemas reg && H.sessionProtocol s == protocolText cfg + then pure () + else throwError (ValidationFailure "Incompatible ReAct checkpoint signature, tool registry, or protocol; explicitly start a new session") + +-- | Start without making a model call. Auto resolves as in the legacy API. +startSession :: (ToPrompt i, ToSchema o, Error ShikumiError :> es) => Signature i o -> ToolRegistry -> ReActConfig -> i -> Eff es ReActSession +startSession sig reg cfg input = do + s <- historyOrThrow (H.newSession (protocolText cfg) (sessionIdentity sig) (sessionToolSchemas reg) (toPrompt input)) + checkSession sig reg cfg s + pure s + +-- | Bind a requested target before the first call. Use a router to supply +-- credentials and model capabilities; checkpoints retain only public identity. +startSessionWithModel :: (ToPrompt i, ToSchema o, Error ShikumiError :> es) => Model -> Signature i o -> ToolRegistry -> ReActConfig -> i -> Eff es ReActSession +startSessionWithModel model sig reg cfg input = do + s <- startSession sig reg cfg input + case C.requestOrigin model of + Nothing -> throwError (ValidationFailure "Explicit session model needs a provider, API, model and credential-free endpoint") + Just origin -> pure (H.bindSessionOrigin (Just origin) s) + +-- | Append a new user turn to a validated checkpoint, resetting its iteration budget. +-- This never dispatches old exchanges. Use 'runSession' or 'advanceSession' next. +continueSession :: (ToPrompt i, ToSchema o, Error ShikumiError :> es) => Signature i o -> ToolRegistry -> ReActConfig -> i -> ReActSession -> Eff es ReActSession +continueSession sig reg cfg input s = do + checkSession sig reg cfg s + historyOrThrow (H.appendUser (toPrompt input) s) + +-- | Make at most one proposal (plus one bounded context retry and summaries). +-- All calls in an accepted proposal are completed sequentially before returning. +advanceSession :: forall i o es. (ToSchema o, FromModel o, Validatable o, LLM :> es, Error ShikumiError :> es) => Signature i o -> ToolRegistry -> ReActConfig -> ReActSession -> Eff es (SessionResult o) +advanceSession sig reg cfg original = do + checkSession sig reg cfg original + if H.sessionFinished original + then throwError (ValidationFailure "Finished session requires a new user turn before advancing") + else + if H.sessionIterations original >= max 0 (maxIters cfg) + then pure (SessionPaused original) + else do + (requestedSession, resp) <- requestRecover original + let resolved = C.requestOrigin (resp ^. #model) + s = H.bindSessionOrigin resolved requestedSession + let payload = resp ^. #message + reject reason = SessionPaused . seal <$> historyOrThrow (H.appendExchange payload [] (Just reason) False s) + if payload ^. #stopReason == B.ErrorReason + then throwError $ case resp ^. #errorInfo of + Just err -> fromBaikaiError err + Nothing -> ProviderFailure (maybe "Model response failed" id (payload ^. #errorMessage)) + else pure () + -- A conforming response echoes the resolved request identity. Never + -- bless old unknown opaque history using a newly selected target. + case H.sessionOrigin requestedSession of + Just expected | resolved /= Just expected -> throwError (ValidationFailure "ReAct response request identity differs from session target") + _ -> pure () + case resolved of + Just _ -> either throwError pure (C.validateReplayOrigin (resp ^. #model) [B.AssistantMessage payload]) + Nothing -> pure () + case parseCalls s resp of + Left reason -> reject reason + Right calls -> case H.validateCalls (priorIds s) calls of + Left (H.HistoryError reason) -> reject reason + Right () + | null calls -> reject "Submit a tool call or shikumi_submit_final with the answer arguments." + | any ((== finalToolName) . (^. #name)) calls -> case calls of + [call] -> case fromModelChecked (call ^. #arguments) of + Left err -> reject ("Invalid final submission: " <> T.pack (show err)) + Right answer -> do + finished <- historyOrThrow (H.appendExchange payload [(call, textToolOutput "Final submission accepted.")] Nothing True s) + pure (SessionFinished answer (seal finished)) + _ -> reject "Final submission cannot be mixed with other calls." + | otherwise -> do + outputs <- + traverse + ( \call -> do + output <- runToolCallOutput reg call + pure (call, either (\err -> ToolOutput (B.toolResultErrorText (renderToolError err)) Nothing []) id output) + ) + calls + next <- historyOrThrow (H.appendExchange payload outputs Nothing False s) + compacted <- + if usageExceedsWindow (compaction cfg) (resp ^. #model) (payload ^. #usage) + then if H.compactionSafe (boundary next) next then forceCompact (seal next) else pure (seal next) + else pure (seal next) + pure (SessionPaused compacted) + where + priorIds s = [c ^. #id_ | H.Exchange _ results Nothing <- H.auditHistory s, (c, _) <- results] + parseCalls s resp + | H.sessionProtocol s == "native" = Right (toolCallsOf resp) + | not (null (toolCallsOf resp)) = Left "Prompt protocol requires JSON text, not native tool calls." + | otherwise = do + actions <- either (\(H.HistoryError t) -> Left t) Right (H.parsePromptActions (responseText resp)) + pure [B.ToolCall ("prompt-" <> T.pack (show (H.sessionTurns s + 1)) <> "-" <> T.pack (show n)) name args | (n, (name, args)) <- zip [1 :: Int ..] actions] + requestContext s = + let native = H.sessionProtocol s == "native" + finalTool = B.emptyTool & #name .~ finalToolName & #description .~ "Submit the final validated answer, alone." & #parameters .~ toSchema (Proxy @o) + sys = + getInstruction sig + <> "\nSubmit your final answer using shikumi_submit_final alone. Its arguments must match: " + <> encodeText (toSchema (Proxy @o)) + <> if native then "" else "\n" <> toolMenu reg <> "\nReply with JSON {\"calls\":[{\"tool\":\"\",\"args\":{...}}]}. Use the same form for shikumi_submit_final." + ctx = buildCtx sys (H.promptMessages s) (if native then registryBaikai reg <> V.singleton finalTool else V.empty) Nothing + in ctx + seal s = H.protectSession (requestContext s) s + request s = do + let ctx = requestContext s + model = maybe emptyModel C.originModel (H.sessionOrigin s) + opts = C.stampContinuation (H.sessionOrigin s) (H.sessionPrefix s) (emptyOptions & #toolChoice .~ (if H.sessionProtocol s == "native" then Just ToolChoiceAuto else Nothing)) + -- Unknown opaque checkpoints fail even for custom interpreters. Resolved + -- target checks also run after routing and at built-in boundaries. + if C.hasOpaqueContinuation (H.promptMessages s) && H.sessionOrigin s == Nothing + then throwError (ValidationFailure "ReAct opaque history has unknown origin; explicitly restart from a caller-approved summary") + else pure () + complete model ctx opts + requestRecover s = + catchError + ((s,) <$> request s) + ( \_cs -> \case + ContextWindowExceeded {} | enabled (compaction cfg) -> do + compacted <- forceCompact s + (compacted,) <$> request compacted + err -> throwError err + ) + boundary s = length (H.auditHistory s) - max 0 (keepRecent (compaction cfg)) + forceCompact s = do + if H.compactionSafe (boundary s) s + then pure () + else historyOrThrow (Left (H.HistoryError "context recovery would change opaque reasoning prefix; explicitly restart from a caller-approved summary")) + let entries = H.promptEntries s + keep = max 0 (keepRecent (compaction cfg)) + through = length (H.auditHistory s) - keep + render = H.renderSummaryEntry + if length entries <= keep || through <= H.sessionCompactedThrough s + then pure s + else do + texts <- compactTail (compaction cfg) emptyModel id id (map render entries) + case texts of + t : _ -> seal <$> historyOrThrow (H.compactSession through t s) + [] -> pure s + +-- | Advance until validated final submission or the per-user-turn iteration limit. +runSession :: (ToSchema o, FromModel o, Validatable o, LLM :> es, Error ShikumiError :> es) => Signature i o -> ToolRegistry -> ReActConfig -> ReActSession -> Eff es (SessionResult o) +runSession sig reg cfg s = do + result <- advanceSession sig reg cfg s + case result of + SessionPaused next | H.sessionIterations next < max 0 (maxIters cfg) -> runSession sig reg cfg next + _ -> pure result diff --git a/src/Shikumi/CodeExec/RLM.hs b/src/Shikumi/CodeExec/RLM.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/CodeExec/RLM.hs @@ -0,0 +1,166 @@ +{-# LANGUAGE ScopedTypeVariables #-} + +-- | Experimental bounded recursive language-model sessions. Documents stay out +-- of the prompt until explicitly inspected; subqueries have depth one. +module Shikumi.CodeExec.RLM + ( RLMConfig (..), + defaultRLMConfig, + RLMOutcome (..), + RLMReport (..), + RLMAudit (..), + rlm, + rlmWithReport, + ) +where + +import Baikai (emptyModel) +import Data.Aeson (Value (..), object, (.=)) +import Data.Proxy (Proxy (..)) +import Data.Text (Text) +import Data.Text qualified as T +import Effectful (Eff, (:>)) +import Effectful.Error.Static (Error, throwError) +import Shikumi.Adapter (ToPrompt (toPrompt), responseText) +import Shikumi.CodeExec.Prompt (encodeText, simpleContext) +import Shikumi.CodeExec.Session +import Shikumi.Error (ShikumiError (..)) +import Shikumi.LLM (LLM, complete) +import Shikumi.Program (Program, embed) +import Shikumi.Schema (FromModel, ToSchema, Validatable, parseOutput, toSchema) +import Shikumi.Signature (Signature, getInstruction) + +data RLMConfig = RLMConfig + { session :: !SessionConfig, + maxOuterIterations :: !Int, + -- | Sum of system and user text characters, checked before every call. + maxRequestChars :: !Int + } + +defaultRLMConfig :: RLMConfig +defaultRLMConfig = RLMConfig defaultSessionConfig 40 64000 + +data RLMOutcome o = RLMSubmitted o | RLMExhausted SessionLimit + deriving stock (Eq, Show) + +data RLMAudit = RLMAudit + { actionText :: !Text, + actionTruncated :: !Bool, + auditObservation :: !(Maybe SessionObservation) + } + deriving stock (Eq, Show) + +data RLMReport = RLMReport + { outerIterations :: !Int, + operations :: !Int, + subqueriesAttempted :: !Int, + totalObservedChars :: !Int, + audit :: ![RLMAudit], + -- | Nothing means typed submission; Just identifies exhaustion. + termination :: !(Maybe SessionLimit) + } + deriving stock (Eq, Show) + +rlm :: (ToPrompt i, ToSchema o, FromModel o, Validatable o) => RLMConfig -> ContextStore -> Signature i o -> Program i o +rlm cfg docs sig = embed $ \i -> do + (out, _) <- runRLM cfg docs sig i + case out of + RLMSubmitted o -> pure o + RLMExhausted limit -> throwError (BudgetExceeded ("RLM session exhausted: " <> T.pack (show limit))) + +rlmWithReport :: (ToPrompt i, ToSchema o, FromModel o, Validatable o) => RLMConfig -> ContextStore -> Signature i o -> Program i (RLMOutcome o, RLMReport) +rlmWithReport cfg docs sig = embed (runRLM cfg docs sig) + +runRLM :: forall i o es. (LLM :> es, Error ShikumiError :> es, ToPrompt i, ToSchema o, FromModel o, Validatable o) => RLMConfig -> ContextStore -> Signature i o -> i -> Eff es (RLMOutcome o, RLMReport) +runRLM cfg docs sig input = do + if maxOuterIterations cfg <= 0 || maxRequestChars cfg <= 0 + then throwError (ValidationFailure "RLM limits must be positive") + else pure () + s <- either (\(SessionError e) -> throwError (ValidationFailure e)) pure (newSession sc docs) + loop 0 s [] "" + where + sc = session cfg + sys = getInstruction sig <> "\n" <> guide <> "\nSubmit value must match: " <> encodeText (toSchema (Proxy @o)) + initial = "Question:\n" <> toPrompt input <> "\nContext metadata:\n" <> encodeText (contextMetadata docs) <> "\nLimits:\n" <> limits + limits = + encodeText $ + object + [ "outerIterations" .= maxOuterIterations cfg, + "requestChars" .= maxRequestChars cfg, + "contextChars" .= maxContextChars sc, + "storedChars" .= maxStoredChars sc, + "actionBytes" .= maxActionBytes sc, + "observationChars" .= maxObservationChars sc, + "observedChars" .= maxObservedChars sc, + "operations" .= maxOperations sc, + "scanChars" .= maxScanChars sc, + "matches" .= maxMatches sc, + "subqueries" .= maxSubqueries sc, + "subqueryChars" .= maxSubqueryChars sc + ] + report n s rows term = RLMReport n (operationCount s) (subqueryAttempts s) (observedChars s) (reverse rows) term + exhausted limit n s rows = pure (RLMExhausted limit, report n s rows (Just limit)) + fits system prompt = toInteger (T.length system) + toInteger (T.length prompt) <= toInteger (maxRequestChars cfg) + loop n s rows history + | n >= maxOuterIterations cfg = exhausted OuterIterations n s rows + | operationCount s >= maxOperations sc = exhausted Operations n s rows + | not (fits sys (initial <> history)) = exhausted RequestCharacters n s rows + | otherwise = do + let (ctx, opts) = simpleContext sys (initial <> history) + raw <- responseText <$> complete emptyModel ctx opts + let parsed = parseSessionAction sc raw + clipped = case parsed of + Right _ -> raw + Left _ -> T.take (maxActionBytes sc `div` 4) raw + row = RLMAudit clipped (clipped /= raw) Nothing + (s1, result) = stepSession s parsed + rows1 = row : rows + history1 = history <> "\nAction: " <> clipped + n1 = n + 1 + case result of + Left limit -> exhausted limit n1 s1 rows1 + Right (Observed obs) -> continue n1 s1 rows1 history1 obs + Right (Submitted value) -> case parseOutput (encodeText value) of + Right o -> pure (RLMSubmitted o, report n1 s1 rows1 Nothing) + Left _ -> continue n1 s1 rows1 history1 (SessionObservation Null (Just (SessionError "submission failed output schema or validation; correct value")) Nothing) + Right (RunQueries prompts) + | any (not . fits subSystem) prompts -> exhausted RequestCharacters n1 s1 rows1 + | otherwise -> queryAll n1 s1 rows1 history1 prompts + continue n s rows history obs = + case retain s rows history obs of + (s1, rows1, _, Left limit) -> exhausted limit n s1 rows1 + (s1, rows1, history1, Right ()) -> loop n s1 rows1 history1 + retain s rows history obs = + let (s1, result) = recordObservation s obs + in case result of + Left limit -> (s1, rows, history, Left limit) + Right bounded -> + let row = RLMAudit "" False (Just bounded) + in (s1, row : rows, history <> "\nObservation: " <> renderObservation bounded, Right ()) + queryAll n s rows history [] = loop n s rows history + queryAll n s rows history (p : ps) = do + let s1 = attemptedSubquery s + (ctx, opts) = simpleContext subSystem p + -- Infrastructure errors propagate unchanged. Reserved slots are never + -- refunded and no further batch member is dispatched after a failure. + answer <- responseText <$> complete (subModel sc) ctx opts + let obs = SessionObservation (String answer) Nothing Nothing + case retain s1 rows history obs of + (s2, rows2, _, Left limit) -> exhausted limit n s2 rows2 + (s2, rows2, history2, Right ()) -> queryAll n s2 rows2 history2 ps + +subSystem :: Text +subSystem = "Answer the supplied question using only its supplied excerpts. You cannot invoke session operations." + +guide :: Text +guide = + T.unlines + [ "Experimental document session. Return exactly one JSON action per turn, no fences.", + "Actions: {op:describe,name}; {op:slice,name,start,count}; {op:find,name,needle,start,limit};", + "{op:store,name,value}; {op:load,name}; {op:query,prompt}; {op:query_batch,prompts:[string]}; {op:submit,value}.", + "Use quoted JSON keys and strings. Offsets/counts are Unicode characters, zero-based, half-open.", + "Find is literal, returns overlapping match offsets and nextOffset for bounded continuation.", + "Documents are immutable; variables persist privately. Names use ASCII letters, digits, underscore or hyphen, at most 128 characters; operation names are reserved.", + "Subqueries run sequentially at depth one; their observations follow prompt order. Include excerpts explicitly.", + "Observations may show marked prefixes. For slices, use nextOffset to continue. Generic truncated values are JSON-text prefixes.", + "Errors consume operations. Correct invalid actions/submissions within the remaining limits. Submit performs no extra model call." + ] diff --git a/src/Shikumi/CodeExec/Session.hs b/src/Shikumi/CodeExec/Session.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/CodeExec/Session.hs @@ -0,0 +1,263 @@ +-- | Experimental, invocation-local data operations. No host execution capability. +module Shikumi.CodeExec.Session + ( ContextStore, + contextStore, + contextMetadata, + SessionConfig (..), + defaultSessionConfig, + validateSessionConfig, + SessionState, + newSession, + operationCount, + subqueryAttempts, + observedChars, + SessionAction (..), + parseSessionAction, + SessionError (..), + SessionLimit (..), + SessionResult (..), + SessionObservation (..), + Truncation (..), + renderObservation, + stepSession, + recordObservation, + attemptedSubquery, + ) +where + +import Baikai (Model, emptyModel) +import Data.Aeson (FromJSON (..), ToJSON (..), Value (..), eitherDecodeStrict, object, withObject, (.:), (.=)) +import Data.ByteString qualified as BS +import Data.Map.Strict (Map) +import Data.Map.Strict qualified as M +import Data.Text (Text) +import Data.Text qualified as T +import Data.Text.Encoding (encodeUtf8) +import Shikumi.CodeExec.Prompt (encodeText) + +newtype ContextStore = ContextStore (Map Text Text) + +-- | Names and contents are immutable. Total size is checked by 'newSession'. +contextStore :: [(Text, Text)] -> Either SessionError ContextStore +contextStore xs + | any (not . validName . fst) xs = Left (SessionError "invalid document name") + | M.size m /= length xs = Left (SessionError "duplicate document name") + | otherwise = Right (ContextStore m) + where + m = M.fromList xs + +contextMetadata :: ContextStore -> Value +contextMetadata (ContextStore m) = object ["documents" .= [object ["name" .= n, "length" .= T.length t] | (n, t) <- M.toList m]] + +data SessionConfig = SessionConfig + { maxContextChars :: !Int, + maxStoredChars :: !Int, + maxActionBytes :: !Int, + maxObservationChars :: !Int, + maxObservedChars :: !Int, + maxOperations :: !Int, + maxScanChars :: !Int, + maxMatches :: !Int, + maxSubqueries :: !Int, + maxSubqueryChars :: !Int, + subModel :: !Model + } + +defaultSessionConfig :: SessionConfig +defaultSessionConfig = SessionConfig 4000000 16000 16000 4000 32000 40 1000000 100 8 8000 emptyModel + +validateSessionConfig :: SessionConfig -> Either SessionError () +validateSessionConfig c + | any (<= 0) [maxContextChars c, maxStoredChars c, maxActionBytes c, maxObservationChars c, maxObservedChars c, maxOperations c, maxScanChars c, maxMatches c, maxSubqueries c, maxSubqueryChars c] = Left (SessionError "all session limits must be positive") + | otherwise = Right () + +data SessionState = SessionState + { sessionConfig :: !SessionConfig, + documents :: !ContextStore, + variables :: !(Map Text Value), + operationCount :: !Int, + reservedSubqueries :: !Int, + subqueryAttempts :: !Int, + observedChars :: !Int + } + +newSession :: SessionConfig -> ContextStore -> Either SessionError SessionState +newSession c docs@(ContextStore m) = do + validateSessionConfig c + if sum [toInteger (T.length n) + toInteger (T.length t) | (n, t) <- M.toList m] > toInteger (maxContextChars c) + then Left (SessionError "context character limit exceeded (including names)") + else Right (SessionState c docs M.empty 0 0 0 0) + +newtype SessionError = SessionError Text deriving stock (Eq, Show) + +data SessionLimit = OuterIterations | Operations | Subqueries | RequestCharacters | ObservationCharacters + deriving stock (Eq, Show) + +data SessionAction + = Describe Text + | Slice Text Int Int + | Find Text Text Int Int + | Store Text Value + | Load Text + | Submit Value + | Query Text + | QueryBatch [Text] + deriving stock (Eq, Show) + +instance ToJSON SessionAction where + toJSON a = case a of + Describe n -> action "describe" ["name" .= n] + Slice n start count -> action "slice" ["name" .= n, "start" .= start, "count" .= count] + Find n needle start limit -> action "find" ["name" .= n, "needle" .= needle, "start" .= start, "limit" .= limit] + Store n v -> action "store" ["name" .= n, "value" .= v] + Load n -> action "load" ["name" .= n] + Submit v -> action "submit" ["value" .= v] + Query p -> action "query" ["prompt" .= p] + QueryBatch ps -> action "query_batch" ["prompts" .= ps] + where + action op fields = object (("op" .= (op :: Text)) : fields) + +instance FromJSON SessionAction where + parseJSON = withObject "session action" $ \o -> do + op <- o .: "op" + case (op :: Text) of + "describe" -> Describe <$> o .: "name" + "slice" -> Slice <$> o .: "name" <*> o .: "start" <*> o .: "count" + "find" -> Find <$> o .: "name" <*> o .: "needle" <*> o .: "start" <*> o .: "limit" + "store" -> Store <$> o .: "name" <*> o .: "value" + "load" -> Load <$> o .: "name" + "submit" -> Submit <$> o .: "value" + "query" -> Query <$> o .: "prompt" + "query_batch" -> QueryBatch <$> o .: "prompts" + _ -> fail "unknown operation" + +parseSessionAction :: SessionConfig -> Text -> Either SessionError SessionAction +parseSessionAction c raw + | BS.length (encodeUtf8 raw) > maxActionBytes c = Left (SessionError "action byte limit exceeded") + | otherwise = case eitherDecodeStrict (encodeUtf8 raw) of + Left _ -> Left (SessionError "invalid action JSON or fields") + Right a -> Right a + +data Truncation = Truncation + { originalChars :: !Int, + displayedChars :: !Int, + nextOffset :: !(Maybe Int) + } + deriving stock (Eq, Show) + +data SessionObservation = SessionObservation + { observationValue :: !Value, + observationError :: !(Maybe SessionError), + truncation :: !(Maybe Truncation) + } + deriving stock (Eq, Show) + +renderObservation :: SessionObservation -> Text +renderObservation o = + encodeText $ + object + [ "value" .= observationValue o, + "error" .= fmap (\(SessionError e) -> e) (observationError o), + "truncation" .= fmap (\t -> object ["originalChars" .= originalChars t, "displayedChars" .= displayedChars t, "nextOffset" .= nextOffset t]) (truncation o) + ] + +data SessionResult = Observed SessionObservation | RunQueries [Text] | Submitted Value + deriving stock (Eq, Show) + +-- | Every action, including a parse error, consumes one operation. Rejected +-- actions preserve variables. Whole batches reserve slots before any dispatch. +stepSession :: SessionState -> Either SessionError SessionAction -> (SessionState, Either SessionLimit SessionResult) +stepSession s action + | operationCount s >= maxOperations c = (s, Left Operations) + | otherwise = case action of + Left e -> bad e + Right a | BS.length (encodeUtf8 (encodeText (toJSON a))) > maxActionBytes c -> bad (SessionError "action byte limit exceeded") + Right a -> case a of + Describe n -> doc n $ \t -> good (object ["name" .= n, "length" .= T.length t]) + Slice n start count -> doc n $ \t -> + if start < 0 || start > T.length t || count < 0 || count > T.length t - start + then bad (SessionError "slice outside document") + else (s1, Right (Observed (sliceObservation n start count t (maxObservationChars c)))) + Find n needle start limit -> doc n $ \t -> + if start < 0 || start > T.length t || limit <= 0 || limit > maxMatches c || T.null needle || T.length needle > maxScanChars c + then bad (SessionError "invalid find offset, needle or match limit") + else + let chunk = T.take (maxScanChars c) (T.drop start t) + (matches, resume) = search needle limit start chunk + end = start + T.length chunk + next = if resume < end then resume else if end < T.length t then max (start + 1) (end - T.length needle + 1) else end + in good (object ["matches" .= matches, "scannedEnd" .= end, "nextOffset" .= next, "complete" .= (next == T.length t)]) + Store n v + | not (validName n) || M.member n dm -> bad (SessionError "reserved or invalid variable name") + | otherwise -> + let vs = M.insert n v (variables s) + in if sum [toInteger (T.length k) + toInteger (T.length (encodeText x)) | (k, x) <- M.toList vs] > toInteger (maxStoredChars c) + then bad (SessionError "stored value character limit exceeded") + else (s1 {variables = vs}, Right (Observed (ok (String "stored")))) + Load n -> maybe (bad (SessionError "unknown variable")) good (M.lookup n (variables s)) + Submit v -> (s1, Right (Submitted v)) + Query p -> queries [p] + QueryBatch ps -> queries ps + where + c = sessionConfig s + ContextStore dm = documents s + s1 = s {operationCount = operationCount s + 1} + bad e = (s1, Right (Observed (SessionObservation Null (Just e) Nothing))) + good v = (s1, Right (Observed (ok v))) + doc n f = maybe (bad (SessionError "unknown document")) f (M.lookup n dm) + queries ps + | null ps || any ((> maxSubqueryChars c) . T.length) ps = bad (SessionError "empty batch or oversized subquery") + | length ps > maxSubqueries c - reservedSubqueries s = (s1, Left Subqueries) + | otherwise = (s1 {reservedSubqueries = reservedSubqueries s + length ps}, Right (RunQueries ps)) + +ok :: Value -> SessionObservation +ok v = SessionObservation v Nothing Nothing + +-- Keep source-relative offsets intact even when JSON escaping expands a slice. +sliceObservation :: Text -> Int -> Int -> Text -> Int -> SessionObservation +sliceObservation name start count source budget = fit (min count budget) + where + fit n = + let v = object ["name" .= name, "start" .= start, "end" .= (start + n), "text" .= T.take n (T.drop start source)] + tr = if n < count then Just (Truncation count n (Just (start + n))) else Nothing + o = SessionObservation v Nothing tr + in if n == 0 || T.length (renderObservation o) <= budget then o else fit (n `div` 2) + +search :: Text -> Int -> Int -> Text -> ([Int], Int) +search needle = go + where + go 0 offset _ = ([], offset) + go remaining offset t = + let (before, after) = T.breakOn needle t + in if T.null after + then ([], offset + T.length t) + else + let found = offset + T.length before + (rest, end) = go (remaining - 1) (found + 1) (T.drop 1 after) + in (found : rest, end) + +-- | Count actual logical LLM dispatch attempts, separately from batch reservations. +attemptedSubquery :: SessionState -> SessionState +attemptedSubquery s = s {subqueryAttempts = subqueryAttempts s + 1} + +-- | Bound the complete serialized observation, including metadata and escaping. +-- Oversized non-slice values are shown as a JSON-text prefix, explicitly marked. +-- A limit too small even for metadata exhausts instead of hiding truncation. +recordObservation :: SessionState -> SessionObservation -> (SessionState, Either SessionLimit SessionObservation) +recordObservation s o + | size > maxObservationChars c || size > maxObservedChars c - observedChars s = (s, Left ObservationCharacters) + | otherwise = (s {observedChars = observedChars s + size}, Right bounded) + where + c = sessionConfig s + raw = encodeText (observationValue o) + bounded + | T.length (renderObservation o) <= maxObservationChars c = o + | truncation o /= Nothing = o + | otherwise = fit (min (T.length raw) (maxObservationChars c)) + fit n = + let candidate = SessionObservation (String (T.take n raw)) (observationError o) (Just (Truncation (T.length raw) n Nothing)) + in if n == 0 || T.length (renderObservation candidate) <= maxObservationChars c then candidate else fit (n `div` 2) + size = T.length (renderObservation bounded) + +validName :: Text -> Bool +validName n = not (T.null n) && T.length n <= 128 && T.all (\x -> x >= 'a' && x <= 'z' || x >= 'A' && x <= 'Z' || x >= '0' && x <= '9' || x == '_' || x == '-') n && n `notElem` ["describe", "slice", "find", "store", "load", "submit", "query", "query_batch"] diff --git a/src/Shikumi/Tool.hs b/src/Shikumi/Tool.hs --- a/src/Shikumi/Tool.hs +++ b/src/Shikumi/Tool.hs @@ -33,6 +33,9 @@ someToolSchema, lowerSomeTool, runErased, + mkDynTool, + runErasedOutput, + runToolCallOutput, ToolRegistry, mkRegistry, registryLookup, @@ -63,9 +66,10 @@ import Data.Vector qualified as V import Effectful (Eff, (:>)) import Effectful.Error.Static (Error, catchError, throwError) -import Shikumi.Error (ShikumiError (..)) +import Shikumi.Error (ShikumiError (..), renderShikumiError) import Shikumi.LLM (LLM) import Shikumi.Schema (FromModel, ToSchema, Validatable, fromModelChecked, toSchema) +import Shikumi.Tool.Output (ToolOutput, renderToolOutput, textToolOutput) -- --------------------------------------------------------------------------- -- The typed tool @@ -117,22 +121,41 @@ (ToSchema i, FromModel i, Validatable i, ToJSON o) => Tool i o -> SomeTool + DynTool :: + Text -> + Text -> + Value -> + (forall es. (LLM :> es, Error ShikumiError :> es) => Value -> Eff es (Either ToolError ToolOutput)) -> + SomeTool +-- | Register a runtime schema and rich-result body without an input Haskell type. +mkDynTool :: + Text -> + Text -> + Value -> + (forall es. (LLM :> es, Error ShikumiError :> es) => Value -> Eff es (Either ToolError ToolOutput)) -> + SomeTool +mkDynTool = DynTool + -- | The name of an erased tool. someToolName :: SomeTool -> Text someToolName (SomeTool t) = name t +someToolName (DynTool n _ _ _) = n -- | The description of an erased tool. someToolDescription :: SomeTool -> Text someToolDescription (SomeTool t) = description t +someToolDescription (DynTool _ d _ _) = d -- | The derived input JSON Schema of an erased tool. someToolSchema :: SomeTool -> Value someToolSchema (SomeTool t) = toolSchemaOf t +someToolSchema (DynTool _ _ s _) = s -- | Lower an erased tool to baikai's wire tool. lowerSomeTool :: SomeTool -> B.Tool lowerSomeTool (SomeTool t) = lowerTool t +lowerSomeTool (DynTool n d s _) = emptyTool & #name .~ n & #description .~ d & #parameters .~ s -- | Run an erased tool against a raw JSON arguments object: decode to the hidden -- @i@, run the body, encode the @o@ to text. A decode failure becomes @@ -144,11 +167,20 @@ SomeTool -> Value -> Eff es (Either ToolError Text) -runErased (SomeTool t) args = +runErased st args = fmap (fmap renderToolOutput) (runErasedOutput st args) + +-- | Rich dispatch; dynamic bodies share typed tools' infrastructure error policy. +runErasedOutput :: (LLM :> es, Error ShikumiError :> es) => SomeTool -> Value -> Eff es (Either ToolError ToolOutput) +runErasedOutput (DynTool n _ _ body) args = + body args `catchError` \_cs e -> + if isInfraToolError e + then throwError e + else pure (Left (ToolRunFailed n (shikumiErrorText e))) +runErasedOutput (SomeTool t) args = case fromModelChecked args of Left err -> pure (Left (ToolArgsInvalid (name t) (shikumiErrorText err))) Right i -> - (Right . encodeText <$> run t i) + (Right . textToolOutput . encodeText <$> run t i) `catchError` \_cs e -> if isInfraToolError e then throwError e @@ -212,6 +244,12 @@ Nothing -> pure (Left (ToolNotFound (tc ^. #name))) Just st -> runErased st (tc ^. #arguments) +-- | Dispatch without projecting away structured result fields. +runToolCallOutput :: (LLM :> es, Error ShikumiError :> es) => ToolRegistry -> ToolCall -> Eff es (Either ToolError ToolOutput) +runToolCallOutput reg tc = case registryLookup (tc ^. #name) reg of + Nothing -> pure (Left (ToolNotFound (tc ^. #name))) + Just st -> runErasedOutput st (tc ^. #arguments) + -- --------------------------------------------------------------------------- -- Helpers -- --------------------------------------------------------------------------- @@ -220,18 +258,9 @@ encodeText :: (ToJSON a) => a -> Text encodeText = decodeUtf8 . LBS.toStrict . encode --- | The human-readable payload of a 'ShikumiError' (its single 'Text' field). +-- | Shared human-readable error rendering for tool observations. shikumiErrorText :: ShikumiError -> Text -shikumiErrorText = \case - InvalidJSON t -> t - MissingField t -> "missing field " <> t - SchemaMismatch t -> t - ValidationFailure t -> t - ProviderFailure t -> t - ContextWindowExceeded t -> t - Timeout t -> t - BudgetExceeded t -> t - CodeExecFailed t -> t +shikumiErrorText = renderShikumiError -- | Which 'ShikumiError's must escape the agent loop rather than become -- observations. Budget and context-window exhaustion are infrastructure faults: diff --git a/src/Shikumi/Tool/Output.hs b/src/Shikumi/Tool/Output.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/Tool/Output.hs @@ -0,0 +1,52 @@ +-- | Lossless local tool results and explicit provider/text projections. +module Shikumi.Tool.Output + ( ToolOutput (..), + textToolOutput, + renderToolOutput, + toolOutputMessage, + ) +where + +import Baikai qualified as B +import Control.Lens ((%~), (&), (^.)) +import Data.Aeson (Value, encode, toJSON) +import Data.ByteString.Lazy qualified as LBS +import Data.Generics.Labels () +import Data.Text (Text) +import Data.Text qualified as T +import Data.Text.Encoding (decodeUtf8) +import Data.Vector qualified as V + +-- | Native content, optional structured data, and uninterpreted extension blocks. +-- Image preservation does not imply that a provider accepts image tool results. +data ToolOutput = ToolOutput + { result :: !B.ToolResult, + structuredContent :: !(Maybe Value), + extensionBlocks :: ![Value] + } + deriving stock (Eq, Show) + +textToolOutput :: Text -> ToolOutput +textToolOutput t = ToolOutput (B.toolResultText t) Nothing [] + +jsonText :: Value -> Text +jsonText = decodeUtf8 . LBS.toStrict . encode + +extraText :: ToolOutput -> [Text] +extraText o = + maybe [] (\v -> ["Structured JSON: " <> jsonText v]) (structuredContent o) + <> map (("Extension JSON: " <>) . jsonText) (extensionBlocks o) + +-- | Text projection preserves image bytes as labeled JSON rather than dropping them. +renderToolOutput :: ToolOutput -> Text +renderToolOutput o = (if result o ^. #isError then "Tool error: " else "") <> T.intercalate "\n" (map render (V.toList (result o ^. #content)) <> extraText o) + where + render (B.ToolResultText (B.TextContent t)) = t + render (B.ToolResultImage img) = "Image JSON: " <> jsonText (toJSON img) + +-- | Preserve native blocks and append explicitly labeled JSON text blocks. +toolOutputMessage :: B.ToolCall -> ToolOutput -> B.Message +toolOutputMessage call o = + B.toolResultFromCall + call + (result o & #content %~ (<> V.fromList (map (B.ToolResultText . B.TextContent) (extraText o)))) diff --git a/test/AcceptanceSpec.hs b/test/AcceptanceSpec.hs --- a/test/AcceptanceSpec.hs +++ b/test/AcceptanceSpec.hs @@ -24,7 +24,6 @@ weatherSignature, weatherTool, ) -import MockLLM (runAgent) import Shikumi.Agent.ReAct ( Action (..), Step (..), @@ -34,6 +33,7 @@ defaultReActConfig, reactWithTrajectory, ) +import Shikumi.Testing (runAgent) import Shikumi.Tool (toolSchemaOf) import Test.Tasty (TestTree, testGroup) import Test.Tasty.HUnit (Assertion, assertBool, assertFailure, testCase, (@?=)) diff --git a/test/AgentHistorySpec.hs b/test/AgentHistorySpec.hs new file mode 100644 --- /dev/null +++ b/test/AgentHistorySpec.hs @@ -0,0 +1,517 @@ +{-# LANGUAGE DataKinds #-} + +module AgentHistorySpec (tests) where + +import Baikai qualified as B +import Baikai.Cost qualified as BC +import Baikai.Error (contentFiltered) +import Baikai.Usage qualified as BU +import Control.Lens ((&), (.~), (^.)) +import Data.Aeson (Value (..), eitherDecode, encode, object, toJSON, (.=)) +import Data.Aeson.Key qualified +import Data.Aeson.KeyMap qualified as KM +import Data.Generics.Labels () +import Data.IORef +import Data.Set qualified as Set +import Data.Text (Text) +import Data.Text qualified as T +import Data.Vector qualified as V +import Effectful (Eff, IOE, liftIO, runEff) +import Effectful.Dispatch.Dynamic (interpret) +import Effectful.Error.Static (Error, runErrorNoCallStack, throwError) +import Fixtures +import GHC.Generics (Generic) +import ReActSessionExample qualified +import Shikumi.Adapter qualified +import Shikumi.Agent.History +import Shikumi.Agent.ReAct +import Shikumi.Compaction (CompactionConfig (..)) +import Shikumi.Error (ShikumiError (..)) +import Shikumi.LLM (LLM (..), complete) +import Shikumi.LLM.Continuation qualified as C +import Shikumi.Schema (FromModel, ToSchema, Validatable (..)) +import Shikumi.Signature (Signature, mkSignature, setInstruction) +import Shikumi.Testing (mkTextResponse, mkToolCallResponse, mkToolCallsResponse) +import Shikumi.Tool +import Shikumi.Tool.Output +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit + +newtype Positive = Positive {value :: Int} + deriving stock (Eq, Show, Generic) + deriving anyclass (ToSchema, FromModel) + +instance Validatable Positive where + validate x@(Positive n) = if n > 0 then Right x else Left "must be positive" + +cfg :: ReActConfig +cfg = defaultReActConfig {protocol = ProtocolNative, compaction = CompactionConfig 0 4 False} + +rich :: ToolOutput +rich = ToolOutput (B.ToolResult (V.singleton (B.ToolResultImage (B.ImageContent "image bytes" "image/png"))) False) (Just (object ["n" .= (42 :: Int)])) [object ["resource" .= String "uri"]] + +registry :: ToolRegistry +registry = mkRegistry [dynamic "A", dynamic "B"] + where + dynamic name = mkDynTool name "test" (object []) $ \_ -> do + _ <- complete B.emptyModel (B.emptyContext & #systemPrompt .~ Just ("dispatch:" <> name)) B.emptyOptions + pure (Right rich) + +firstTurn :: B.Response +firstTurn = mkToolCallsResponse [("call-A", "A", object []), ("call-B", "B", object [])] + +finalTurn :: B.Response +finalTurn = mkToolCallResponse "final-1" finalToolName (toJSON expectedWeather) + +-- The test interpreter observes dispatch via a nested LLM operation, respecting +-- tool bodies' rank-polymorphic effect boundary without unsafe IO in a tool. +recording :: [Either ShikumiError B.Response] -> Eff '[LLM, Error ShikumiError, IOE] a -> IO (Either ShikumiError a, [B.Context], [Text]) +recording script action = do + scriptRef <- newIORef script + contexts <- newIORef [] + dispatches <- newIORef [] + result <- + runEff + . runErrorNoCallStack + . interpret + ( \_ -> \case + Complete model ctx opts -> case ctx ^. #systemPrompt of + Just sys | "dispatch:" `T.isPrefixOf` sys -> do + liftIO (modifyIORef' dispatches (<> [T.drop 9 sys])) + pure (mkTextResponse "done") + _ -> do + either throwError pure (C.validateRequestContinuation model ctx opts) + liftIO (modifyIORef' contexts (<> [ctx])) + next <- + liftIO + ( atomicModifyIORef' + scriptRef + ( \xs -> case xs of + [] -> ([], Left (ProviderFailure "unexpected model call")) + x : rest -> (rest, x) + ) + ) + either throwError pure next + Stream {} -> pure [] + ) + $ action + (result,,) <$> readIORef contexts <*> readIORef dispatches + +start :: Eff '[LLM, Error ShikumiError, IOE] ReActSession +start = startSession weatherSignature registry cfg weatherQuestion + +paused :: SessionResult o -> Eff '[LLM, Error ShikumiError, IOE] ReActSession +paused (SessionPaused s) = pure s +paused _ = throwError (ValidationFailure "expected paused session") + +restored :: ReActSession -> Eff '[LLM, Error ShikumiError, IOE] ReActSession +restored s = either (throwError . ValidationFailure . T.pack . show) pure (decodeSession (encodeSession s)) + +assertFinished :: Either ShikumiError (SessionResult WeatherResp) -> Assertion +assertFinished (Right (SessionFinished answer _)) = answer @?= expectedWeather +assertFinished other = assertFailure (show other) + +mutate :: Text -> Value -> Value -> Value +mutate key v (Object o) = Object (KM.insert (fromStringKey key) v o) + where + fromStringKey = Data.Aeson.Key.fromText +mutate _ _ v = v + +tests :: TestTree +tests = + testGroup + "AgentHistory" + [ testCase "documented continuation example" $ do + ReActSessionExample.example >>= (@?= Right ("Paris, France", 2)), + testCase "native two-tool checkpoint resumes with exact messages and no extraction" $ do + (answer, requests, dispatched) <- recording (map Right [firstTurn, finalTurn]) $ do + s <- start >>= advanceSession weatherSignature registry cfg >>= paused + copy <- restored s + if copy == s then pure () else throwError (ValidationFailure "round-trip differs") + resumed <- continueSession weatherSignature registry cfg weatherQuestion copy + advanceSession weatherSignature registry cfg resumed + assertFinished answer + dispatched @?= ["A", "B"] + length requests @?= 2 + case requests of + [_, second] -> + second ^. #messages + @?= V.fromList + [ B.user (Shikumi.Adapter.toPrompt weatherQuestion), + B.AssistantMessage (firstTurn ^. #message), + toolOutputMessage (B.ToolCall "call-A" "A" (object [])) rich, + toolOutputMessage (B.ToolCall "call-B" "B" (object [])) rich, + B.user (Shikumi.Adapter.toPrompt weatherQuestion) + ] + _ -> assertFailure "wrong request count", + testCase "resumed and uninterrupted next requests are identical" $ do + let execute restore = recording (map Right [firstTurn, finalTurn]) $ do + s <- start >>= advanceSession weatherSignature registry cfg >>= paused + copy <- if restore then restored s else pure s + next <- continueSession weatherSignature registry cfg weatherQuestion copy + advanceSession weatherSignature registry cfg next + (_, uninterrupted, _) <- execute False + (_, resumed, _) <- execute True + resumed @?= uninterrupted, + testCase "message metadata, thinking signature and rich blocks survive JSON bytes" $ do + let replay = B.ThinkingReplay B.OpenAIResponses "test-model" (V.singleton (object ["opaque" .= ("replay" :: String)])) + thinking = B.ThinkingContent "opaque" (Just "signature") True (Just replay) + response = firstTurn & #message . #content .~ V.cons (B.AssistantThinking thinking) (firstTurn ^. #message . #content) + (result, _, _) <- recording [Right response] (start >>= advanceSession weatherSignature registry cfg >>= paused) + case result of + Right s -> (eitherDecode (encode (encodeSession s)) >>= either (Left . show) Right . decodeSession) @?= Right s + Left e -> assertFailure (show e), + testCase "separate Claude and Responses continuations survive resume exactly" $ do + mapM_ + ( \(model, thinking) -> do + let response = firstTurn & #model .~ model & #message . #content .~ V.cons (B.AssistantThinking thinking) (firstTurn ^. #message . #content) + (answer, requests, dispatched) <- recording [Right response, Right (finalTurn & #model .~ model)] $ do + s <- startSessionWithModel model weatherSignature registry cfg weatherQuestion >>= advanceSession weatherSignature registry cfg >>= paused >>= restored + advanceSession weatherSignature registry cfg s + assertFinished answer + dispatched @?= ["A", "B"] + case requests of + [_, ctx] -> assertBool "exact assistant" (B.AssistantMessage (response ^. #message) `elem` V.toList (ctx ^. #messages)) + _ -> assertFailure "wrong requests" + ) + reasoningFixtures, + testCase "safe summary excludes opaque sentinels while audit preserves them" $ do + let thinking = B.ThinkingContent "SECRET-REDACTED" (Just "SECRET-SIGNATURE") True (Just (B.ThinkingReplay B.OpenAIResponses "m" (V.singleton (String "SECRET-REPLAY")))) + payload = firstTurn ^. #message & #content .~ V.cons (B.AssistantThinking thinking) (firstTurn ^. #message . #content) + built = newSession "native" (object []) (object []) "question" >>= appendExchange payload [(B.ToolCall "call-A" "A" (object []), rich), (B.ToolCall "call-B" "B" (object []), rich)] Nothing False + case built of + Left e -> assertFailure (show e) + Right s -> do + let rendered = renderSessionSummaryInput s + audit = show (encodeSession s) + mapM_ + ( \secret -> do + assertBool "absent from summary" (not (secret `T.isInfixOf` rendered)) + assertBool "present in audit" (secret `T.isInfixOf` T.pack audit) + ) + ["SECRET-REDACTED", "SECRET-SIGNATURE", "SECRET-REPLAY"] + assertBool "tool names retained" ("A" `T.isInfixOf` rendered) + decodeSession (encodeSession s) @?= Right s, + testCase "opaque compaction defers proactively and overflow makes no summary call" $ do + let (model, thinking) = claudeFixture + compactCfg = cfg {compaction = CompactionConfig 0 1 True} + response = firstTurn & #model .~ (model & #contextWindow .~ 10) & #message . #usage . #inputTokens .~ 10 & #message . #content .~ V.cons (B.AssistantThinking thinking) (firstTurn ^. #message . #content) + (result, requests, dispatched) <- recording [Right response, Left (ContextWindowExceeded "full")] $ do + s <- startSessionWithModel model weatherSignature registry compactCfg weatherQuestion >>= advanceSession weatherSignature registry compactCfg >>= paused + if sessionCompactedThrough s == 0 then pure () else throwError (ValidationFailure "unexpected compaction") + advanceSession weatherSignature registry compactCfg s + case result of + Left (ValidationFailure message) -> assertBool "actionable" ("restart" `T.isInfixOf` message) + _ -> assertFailure (show result) + length requests @?= 2 + dispatched @?= ["A", "B"], + testCase "automatic summary requests contain no discarded opaque payload" $ do + mapM_ + ( \(model, thinking) -> do + let compactCfg = cfg {compaction = CompactionConfig 0 1 True} + response = firstTurn & #model .~ model & #message . #content .~ V.cons (B.AssistantThinking thinking) (firstTurn ^. #message . #content) + later = mkToolCallResponse "call-C" "A" (object []) & #model .~ (model & #contextWindow .~ 10) & #message . #usage . #inputTokens .~ 10 + (result, requests, dispatched) <- recording [Right response, Right later, Right (mkTextResponse "safe summary")] $ do + s <- startSessionWithModel model weatherSignature registry compactCfg weatherQuestion >>= advanceSession weatherSignature registry compactCfg >>= paused + advanceSession weatherSignature registry compactCfg s >>= paused + case result of + Left e -> assertFailure (show e) + Right s -> do + sessionCompactedThrough s @?= 2 + length (auditHistory s) @?= 3 + dispatched @?= ["A", "B", "A"] + case reverse requests of + summaryRequest : _ -> do + let rendered = T.pack (show (toJSON summaryRequest)) + mapM_ (\secret -> assertBool "opaque absent from actual summary request" (not (secret `T.isInfixOf` rendered))) ["signature-bytes", "encrypted-bytes", "redacted-secret"] + _ -> assertFailure "no summary request" + length requests @?= 3 + ) + (reasoningFixtures <> [(B.mkModel B.AnthropicMessages "claude-test" "https://claude.example", B.ThinkingContent "redacted-secret" Nothing True Nothing)]), + testCase "response request mismatch rejects before tool dispatch" $ do + let (model, _) = claudeFixture + (result, requests, dispatched) <- recording [Right (firstTurn & #model .~ (model & #modelId .~ "changed"))] $ do + s <- startSessionWithModel model weatherSignature registry cfg weatherQuestion + advanceSession weatherSignature registry cfg s + assertBool "mismatch rejected" (isLeft result) + length requests @?= 1 + dispatched @?= [], + testCase "legacy opaque origin is unknown; explicit restart leaves audit unchanged" $ do + let (model, thinking) = claudeFixture + response = firstTurn & #model .~ model & #message . #content .~ V.cons (B.AssistantThinking thinking) (firstTurn ^. #message . #content) + (built, _, _) <- recording [Right response] (startSessionWithModel model weatherSignature registry cfg weatherQuestion >>= advanceSession weatherSignature registry cfg >>= paused) + case built of + Left e -> assertFailure (show e) + Right old -> do + let legacy = mutate "version" (Number 1) (encodeSession old) + case decodeSession legacy of + Left e -> assertFailure (show e) + Right unknown -> do + sessionOrigin unknown @?= Nothing + (rejected, requests, dispatched) <- recording [] (advanceSession weatherSignature registry cfg unknown) + assertBool "legacy rejected" (isLeft rejected) + requests @?= [] + dispatched @?= [] + case restartSessionFromSummary "Caller approved summary" unknown of + Left e -> assertFailure (show e) + Right fresh -> do + sessionOrigin fresh @?= Nothing + sessionTurns fresh @?= 0 + length (auditHistory fresh) @?= 1 + assertBool "no native old messages" (not (C.hasOpaqueContinuation (promptMessages fresh))) + length (auditHistory unknown) @?= 2 + decodeSession (encodeSession old) @?= Right old + (answer, freshRequests, freshDispatch) <- recording [Right finalTurn] (advanceSession weatherSignature registry cfg fresh) + assertFinished answer + length freshRequests @?= 1 + freshDispatch @?= [], + testCase "immutable version-one checkpoint reads and upgrades losslessly" $ do + let bytes = "{\"version\":1,\"protocol\":\"prompt\",\"fingerprint\":{},\"tools\":{},\"turns\":0,\"iterations\":0,\"finished\":false,\"history\":[{\"kind\":\"user\",\"content\":[{\"type\":\"user_text\",\"data\":{\"text\":\"question\"}}],\"timestamp\":null}],\"compactedThrough\":0,\"summary\":null}" + (eitherDecode bytes >>= either (Left . show) Right . decodeSession) @?= either (Left . show) Right (newSession "prompt" (object []) (object []) "question"), + testCase "version-one plain session binds only after a resolved response" $ do + (result, _, _) <- recording [] start + case result of + Left e -> assertFailure (show e) + Right initial -> case decodeSession (mutate "version" (Number 1) (encodeSession initial)) of + Left e -> assertFailure (show e) + Right legacy -> do + sessionOrigin legacy @?= Nothing + let (model, _) = claudeFixture + (answer, _, _) <- recording [Right (firstTurn & #model .~ model)] (advanceSession weatherSignature registry cfg legacy >>= paused) + case answer of + Right bound -> sessionOrigin bound @?= C.requestOrigin model + Left e -> assertFailure (show e), + testCase "all invalid native proposals are audit-only and dispatch nothing" $ do + let invalid = + [ mkToolCallsResponse [("", "A", object []), ("ok", "B", object [])], + mkToolCallsResponse [("same", "A", object []), ("same", "B", object [])], + mkToolCallsResponse [("ok", "A", object []), ("cut", "B", String "{")], + mkToolCallsResponse [("a", "A", object []), ("f", finalToolName, toJSON expectedWeather)], + mkToolCallResponse "bad-final" finalToolName (object []) + ] + mapM_ + ( \response -> do + (result, requests, dispatched) <- recording [Right response, Right finalTurn] $ do + s <- start >>= advanceSession weatherSignature registry cfg >>= paused + copy <- restored s + advanceSession weatherSignature registry cfg copy + assertFinished result + dispatched @?= [] + case requests of + [_, ctx] -> do + assertBool "invalid assistant omitted" (not (B.AssistantMessage (response ^. #message) `elem` V.toList (ctx ^. #messages))) + assertBool "no orphan tool results" (null [() | B.ToolResultMessage _ <- V.toList (ctx ^. #messages)]) + _ -> assertFailure "wrong requests" + ) + invalid, + testCase "IDs cannot be reused in later accepted exchanges" $ do + (result, _, dispatched) <- recording (map Right [firstTurn, firstTurn, finalTurn]) $ do + s <- start >>= advanceSession weatherSignature registry cfg >>= paused + next <- advanceSession weatherSignature registry cfg s >>= paused + advanceSession weatherSignature registry cfg next + assertFinished result + dispatched @?= ["A", "B"], + testCase "unknown versions, unresolved exchanges and corrupt JSON are rejected" $ do + (result, _, _) <- recording [Right firstTurn] (start >>= advanceSession weatherSignature registry cfg >>= paused) + case result of + Right s -> do + assertBool "version" (isLeft (decodeSession (mutate "version" (Number 99) (encodeSession s)))) + let corruptResults (Object o) = case KM.lookup "history" o of + Just (Array entries) -> + Object + ( KM.insert + "history" + ( Array + ( V.map + ( \e -> case e of + Object fields | KM.lookup "kind" fields == Just (String "exchange") -> Object (KM.insert "results" (Array V.empty) fields) + _ -> e + ) + entries + ) + ) + o + ) + _ -> Object o + corruptResults v = v + assertBool "unresolved" (isLeft (decodeSession (corruptResults (encodeSession s)))) + assertBool "invalid JSON" (isLeft (eitherDecode "{" :: Either String Value)) + Left e -> assertFailure (show e), + testCase "signature and registry mismatch fail before any execution" $ do + (result, requests, dispatched) <- recording [] $ do + s <- start + advanceSession (setInstruction "changed" weatherSignature) registry cfg s + assertBool "signature rejected" (isLeft result) + requests @?= [] + dispatched @?= [] + (changed, reqs, _) <- recording [] $ do + s <- start + advanceSession weatherSignature (mkRegistry []) cfg s + assertBool "registry rejected" (isLeft changed) + reqs @?= [], + testCase "reserved final tool collision is rejected at startup" $ do + (result, requests, _) <- recording [] (startSession weatherSignature (mkRegistry [mkDynTool finalToolName "" (object []) (\_ -> pure (Right rich))]) cfg weatherQuestion) + assertBool "collision" (isLeft result) + requests @?= [], + testCase "iteration exhaustion returns a checkpoint without another call" $ do + let limited = cfg {maxIters = 1} + (result, requests, dispatched) <- recording [Right firstTurn] $ do + s <- startSession weatherSignature registry limited weatherQuestion + runSession weatherSignature registry limited s + case result of + Right (SessionPaused s) -> sessionTurns s @?= 1 + _ -> assertFailure (show result) + length requests @?= 1 + dispatched @?= ["A", "B"], + testCase "prompt fallback preserves continuation and synthetic IDs" $ do + let promptCfg = cfg {protocol = ProtocolPrompt} + proposal = mkTextResponse "{\"calls\":[{\"tool\":\"A\",\"args\":{}},{\"tool\":\"B\",\"args\":{}}]}" + final = mkTextResponse "{\"calls\":[{\"tool\":\"shikumi_submit_final\",\"args\":{\"tempC\":12,\"summary\":\"mild\"}}]}" + (result, requests, dispatched) <- recording (map Right [proposal, final]) $ do + s <- startSession weatherSignature registry promptCfg weatherQuestion >>= advanceSession weatherSignature registry promptCfg >>= paused >>= restored + next <- continueSession weatherSignature registry promptCfg weatherQuestion s + advanceSession weatherSignature registry promptCfg next + assertFinished result + dispatched @?= ["A", "B"] + length requests @?= 2, + testCase "context retry compacts whole exchanges without redispatch" $ do + let compactCfg = cfg {compaction = CompactionConfig 0 1 True} + (result, requests, dispatched) <- recording [Right firstTurn, Left (ContextWindowExceeded "full"), Right (mkTextResponse "Earlier question"), Right finalTurn] $ do + s <- startSession weatherSignature registry compactCfg weatherQuestion >>= advanceSession weatherSignature registry compactCfg >>= paused + advanceSession weatherSignature registry compactCfg s + assertFinished result + dispatched @?= ["A", "B"] + length requests @?= 4 + case result of + Right (SessionFinished _ s) -> do + decodeSession (encodeSession s) @?= Right s + length (auditHistory s) @?= 3 + _ -> pure (), + testCase "usage compaction retains complete exchanges across restore" $ do + let compactCfg = cfg {compaction = CompactionConfig 0 1 True} + large = firstTurn & #model . #contextWindow .~ 10 & #message . #usage . #inputTokens .~ 10 + (result, requests, dispatched) <- recording (map Right [large, mkTextResponse "question summary", finalTurn]) $ do + s <- startSession weatherSignature registry compactCfg weatherQuestion >>= advanceSession weatherSignature registry compactCfg >>= paused >>= restored + advanceSession weatherSignature registry compactCfg s + assertFinished result + dispatched @?= ["A", "B"] + length requests @?= 3 + case reverse requests of + ctx : _ -> do + length [() | B.AssistantMessage _ <- V.toList (ctx ^. #messages)] @?= 1 + length [() | B.ToolResultMessage _ <- V.toList (ctx ^. #messages)] @?= 2 + _ -> assertFailure "missing request", + testCase "tool infrastructure failure aborts remaining dispatch without retry" $ do + let failing = mkRegistry [mkDynTool "A" "" (object []) (\_ -> throwError (BudgetExceeded "tool budget")), mkDynTool "B" "" (object []) (\_ -> pure (Right rich))] + (result, requests, dispatched) <- recording [Right firstTurn] $ do + s <- startSession weatherSignature failing cfg weatherQuestion + advanceSession weatherSignature failing cfg s + result @?= Left (BudgetExceeded "tool budget") + length requests @?= 1 + dispatched @?= [], + testCase "provider error payload is never dispatched" $ do + let failed = firstTurn & #message . #stopReason .~ B.ErrorReason & #message . #errorMessage .~ Just "transport failed" + (result, _, dispatched) <- recording [Right failed] (start >>= advanceSession weatherSignature registry cfg) + result @?= Left (ProviderFailure "transport failed") + dispatched @?= [], + testCase "structured refusal preserves error and never dispatches partial calls" $ do + let err = contentFiltered "refused" + failed = firstTurn & #message . #stopReason .~ B.ErrorReason & #errorInfo .~ Just err + (result, requests, dispatched) <- recording [Right failed] (start >>= advanceSession weatherSignature registry cfg) + result @?= Left (ProviderError err) + length requests @?= 1 + dispatched @?= [], + testCase "ordinary tool failures are error-flagged model messages" $ do + let unknown = mkToolCallResponse "unknown-id" "unknown-tool" (object []) + (result, requests, _) <- recording (map Right [unknown, finalTurn]) $ do + s <- start >>= advanceSession weatherSignature registry cfg >>= paused + advanceSession weatherSignature registry cfg s + assertFinished result + case reverse requests of + ctx : _ -> [p ^. #isError | B.ToolResultMessage p <- V.toList (ctx ^. #messages)] @?= [True] + _ -> assertFailure "missing request", + testCase "checkpoint rejects altered result names and duplicate results" $ do + (result, _, _) <- recording [Right firstTurn] (start >>= advanceSession weatherSignature registry cfg >>= paused) + case result of + Right s -> case auditHistory s of + [_, Exchange payload results Nothing] -> do + let base = newSession "native" (sessionFingerprint s) (sessionTools s) "question" + changed = case results of + (call, output) : rest -> (call & #name .~ "different", output) : rest + [] -> [] + assertBool "name mismatch" (isLeft (base >>= appendExchange payload changed Nothing False)) + assertBool "duplicate result" (isLeft (base >>= appendExchange payload (results <> results) Nothing False)) + _ -> assertFailure "unexpected audit" + Left err -> assertFailure (show err), + testCase "schema changes under an existing tool name require restart" $ do + let changed = mkRegistry [mkDynTool "A" "test" (object ["type" .= String "string"]) (\_ -> pure (Right rich)), mkDynTool "B" "test" (object []) (\_ -> pure (Right rich))] + (result, requests, _) <- recording [] $ do + s <- start + continueSession weatherSignature changed cfg weatherQuestion s + assertBool "schema rejected" (isLeft result) + requests @?= [], + testCase "fenced prompt uses identical dispatch and persistence parsing" $ do + let promptCfg = cfg {protocol = ProtocolPrompt} + proposal = mkTextResponse "```json\n{\"calls\":[{\"tool\":\"A\",\"args\":{}}]}" + (result, _, dispatched) <- recording [Right proposal] $ do + s <- startSession weatherSignature registry promptCfg weatherQuestion + advanceSession weatherSignature registry promptCfg s >>= paused >>= restored + assertBool "checkpoint returned" (not (isLeft result)) + dispatched @?= ["A"], + testCase "final submission runs custom semantic validation" $ do + let sig = mkSignature "Return a positive value" :: Signature AnswerWeatherQuestion Positive + bad = mkToolCallResponse "bad" finalToolName (object ["value" .= (-1 :: Int)]) + good = mkToolCallResponse "good" finalToolName (object ["value" .= (2 :: Int)]) + (result, requests, dispatched) <- recording (map Right [bad, good]) $ do + s <- startSession sig registry cfg weatherQuestion + runSession sig registry cfg s + case result of + Right (SessionFinished answer s) -> do + answer @?= Positive 2 + sessionTurns s @?= 2 + _ -> assertFailure (show result) + length requests @?= 2 + dispatched @?= [], + testCase "exact rational costs and timestamps survive checkpoint encoding" $ do + let cost = BC.Cost (1 / 3) (BC.CostBreakdown (1 / 7) (2 / 9) (1 / 11) (1 / 13)) (BC.CostBasis (Set.singleton BC.StandardTokenRates) (Set.singleton BC.CacheWriteUsageNotReported)) + response = + firstTurn + & #message . #usage . #cost .~ cost + & #message . #usage . #reasoningTokens .~ Just 5 + & #message . #usage . #availability .~ Just (BU.UsageAvailability (Set.singleton BU.CacheWriteUsage) False (Set.singleton (BU.BillingServiceTier "standard"))) + & #message . #timestamp .~ Just (read "2026-09-07 01:00:00 UTC") + & #message . #errorMessage .~ Just "nonfatal diagnostic" + (result, _, _) <- recording [Right response] (start >>= advanceSession weatherSignature registry cfg >>= paused) + case result of + Right s -> (eitherDecode (encode (encodeSession s)) >>= either (Left . show) Right . decodeSession) @?= Right s + Left err -> assertFailure (show err), + testCase "legacy checkpoints without billing metadata still decode" $ do + let stripBilling (Object fields) = Object (KM.mapWithKey (\_ -> stripBilling) (KM.delete "basis" (KM.delete "availability" fields))) + stripBilling (Array values) = Array (V.map stripBilling values) + stripBilling other = other + (result, _, _) <- recording [Right firstTurn] (start >>= advanceSession weatherSignature registry cfg >>= paused) + case result of + Right s -> decodeSession (stripBilling (encodeSession s)) @?= Right s + Left err -> assertFailure (show err), + testCase "context retry is bounded" $ do + let compactCfg = cfg {compaction = CompactionConfig 0 1 True} + (result, requests, dispatched) <- recording [Right firstTurn, Left (ContextWindowExceeded "full"), Right (mkTextResponse "summary"), Left (ContextWindowExceeded "again")] $ do + s <- startSession weatherSignature registry compactCfg weatherQuestion >>= advanceSession weatherSignature registry compactCfg >>= paused + advanceSession weatherSignature registry compactCfg s + result @?= Left (ContextWindowExceeded "again") + length requests @?= 4 + dispatched @?= ["A", "B"] + ] + where + isLeft (Left _) = True + isLeft _ = False + +-- Separate API-valid shapes; neither mixes Claude signatures and Responses items. +reasoningFixtures :: [(B.Model, B.ThinkingContent)] +reasoningFixtures = + [ claudeFixture, + (B.mkModel B.OpenAIResponses "responses-test" "https://responses.example", B.ThinkingContent "" Nothing False (Just (B.ThinkingReplay B.OpenAIResponses "responses-test" (V.singleton (object ["type" .= String "reasoning", "id" .= String "rs_1", "encrypted_content" .= String "encrypted-bytes", "summary" .= ([] :: [Value])]))))) + ] + +claudeFixture :: (B.Model, B.ThinkingContent) +claudeFixture = (B.mkModel B.AnthropicMessages "claude-test" "https://claude.example", B.ThinkingContent "" (Just "signature-bytes") False Nothing) diff --git a/test/BuiltinAcceptanceSpec.hs b/test/BuiltinAcceptanceSpec.hs --- a/test/BuiltinAcceptanceSpec.hs +++ b/test/BuiltinAcceptanceSpec.hs @@ -14,7 +14,6 @@ import Data.Text.IO qualified as TIO import Data.Vector qualified as V import GHC.Generics (Generic) -import MockLLM (mkTextResponse, mkToolCallResponse, runAgent) import Shikumi.Adapter (ToPrompt) import Shikumi.Agent.ReAct ( Action (..), @@ -27,6 +26,7 @@ ) import Shikumi.Schema (FromModel, ToSchema, Validatable) import Shikumi.Signature (Signature, mkSignature) +import Shikumi.Testing (mkTextResponse, mkToolCallResponse, runAgent) import Shikumi.Tool.Builtin (builtinRegistry) import Shikumi.Tool.Env (localToolEnv) import Shikumi.Tool.Web (FetchResult (..), SearchResult (..), WebClient (..)) diff --git a/test/CodeActSpec.hs b/test/CodeActSpec.hs --- a/test/CodeActSpec.hs +++ b/test/CodeActSpec.hs @@ -18,7 +18,6 @@ import Effectful (runEff) import Effectful.Error.Static (runErrorNoCallStack) import GHC.Generics (Generic) -import MockLLM (mkTextResponse, mkUsageResponse, runAgent, runMockLLMThrowingOn) import Shikumi.Adapter (ToPrompt) import Shikumi.Agent.ReAct (Action (..), Step (..), Termination (..), Trajectory (..)) import Shikumi.CodeExec.CodeAct (CodeActConfig (..), codeActWithTrajectory, defaultCodeActConfig) @@ -29,6 +28,7 @@ import Shikumi.Program (runProgram) import Shikumi.Schema (FromModel, ToSchema, Validatable) import Shikumi.Signature (Signature, mkSignature) +import Shikumi.Testing (mkTextResponse, mkUsageResponse, runAgent, runScriptLLMThrowingOn) import Shikumi.Tool (SomeTool (..), Tool, ToolRegistry, mkRegistry, mkTool) import Test.Tasty (TestTree, testGroup) import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=)) @@ -138,7 +138,7 @@ out <- runEff . runErrorNoCallStack @ShikumiError - . runMockLLMThrowingOn [2] (ContextWindowExceeded "context length exceeded") script + . runScriptLLMThrowingOn [2] (ContextWindowExceeded "context length exceeded") script $ runProgram prog (Task "compute 42") case out of Right (answer, traj) -> do @@ -161,7 +161,7 @@ out <- runEff . runErrorNoCallStack @ShikumiError - . runMockLLMThrowingOn [2] (ContextWindowExceeded "context length exceeded") script + . runScriptLLMThrowingOn [2] (ContextWindowExceeded "context length exceeded") script $ runProgram prog (Task "compute 42") out @?= Left (ContextWindowExceeded "context length exceeded") ] diff --git a/test/CompactionSpec.hs b/test/CompactionSpec.hs --- a/test/CompactionSpec.hs +++ b/test/CompactionSpec.hs @@ -1,9 +1,11 @@ module CompactionSpec (tests) where import Baikai (emptyModel, zeroUsage) -import Control.Lens ((&), (.~)) +import Baikai.Speed (Speed (..)) +import Control.Lens ((&), (.~), (^.)) import Data.Aeson (Value (..)) import Data.Generics.Labels () +import Data.IORef (newIORef, readIORef) import Data.Text (Text) import Data.Vector qualified as V import Effectful (runEff) @@ -15,13 +17,6 @@ weatherRegistry, weatherSignature, ) -import MockLLM - ( mkTextResponse, - mkUsageResponse, - runAgent, - runEffMock, - runMockLLMThrowingOn, - ) import Shikumi.Agent.ReAct ( Action (..), ReActConfig (..), @@ -40,7 +35,16 @@ usageExceedsWindow, ) import Shikumi.Error (ShikumiError (..)) +import Shikumi.LLM.Defaults import Shikumi.Program (runProgram) +import Shikumi.Testing + ( captureLLMRequests, + mkTextResponse, + mkUsageResponse, + runEffScript, + runScriptLLM, + runScriptLLMThrowingOn, + ) import Test.Tasty (TestTree, testGroup) import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=)) @@ -61,13 +65,13 @@ testCase "compactTail folds older items and keeps the recent tail" $ do let cfg = defaultCompactionConfig {keepRecent = 2} res <- - runEffMock [mkTextResponse "S"] $ + runEffScript [mkTextResponse "S"] $ compactTail cfg emptyModel id ("summary:" <>) (["e1", "e2", "e3", "e4", "e5", "e6"] :: [Text]) res @?= Right ["summary:S", "e5", "e6"], testCase "compactTail with enabled=False is the identity and calls no model" $ do let cfg = defaultCompactionConfig {enabled = False, keepRecent = 0} items = ["e1", "e2", "e3"] :: [Text] - res <- runEffMock [] $ compactTail cfg emptyModel id ("summary:" <>) items + res <- runEffScript [] $ compactTail cfg emptyModel id ("summary:" <>) items res @?= Right items, testCase "agent on tiny window compacts and completes" $ do let cfg = @@ -84,7 +88,17 @@ mkUsageResponse model 10 finishReply, mkTextResponse extractReply ] - res <- runAgent script (reactWithTrajectory weatherSignature weatherRegistry cfg) weatherQuestion + captured <- newIORef [] + res <- + runEff + . runErrorNoCallStack @ShikumiError + . runScriptLLM script + . captureLLMRequests captured + . withRequestDefaults (emptyRequestDefaults {defaultSpeed = Just SpeedFast}) + $ runProgram (reactWithTrajectory weatherSignature weatherRegistry cfg) weatherQuestion + calls <- readIORef captured + length calls @?= 5 + map (\(_, _, opts) -> opts ^. #speed) calls @?= replicate 5 (Just SpeedFast) case res of Right (o :: WeatherResp, traj) -> do o @?= expectedWeather @@ -119,7 +133,7 @@ res <- runEff . runErrorNoCallStack @ShikumiError - . runMockLLMThrowingOn [2] (ContextWindowExceeded "context length exceeded") script + . runScriptLLMThrowingOn [2] (ContextWindowExceeded "context length exceeded") script $ runProgram prog weatherQuestion case res of Right (_ :: WeatherResp, traj) -> @@ -142,7 +156,7 @@ res <- runEff . runErrorNoCallStack @ShikumiError - . runMockLLMThrowingOn [2] (ContextWindowExceeded "context length exceeded") script + . runScriptLLMThrowingOn [2] (ContextWindowExceeded "context length exceeded") script $ runProgram prog weatherQuestion res @?= Left (ContextWindowExceeded "context length exceeded"), testCase "extract overflow is caught, compacted, and retried once" $ do @@ -163,7 +177,7 @@ res <- runEff . runErrorNoCallStack @ShikumiError - . runMockLLMThrowingOn [3] (ContextWindowExceeded "context length exceeded") script + . runScriptLLMThrowingOn [3] (ContextWindowExceeded "context length exceeded") script $ runProgram prog weatherQuestion case res of Right (_ :: WeatherResp, traj) -> @@ -185,7 +199,7 @@ res <- runEff . runErrorNoCallStack @ShikumiError - . runMockLLMThrowingOn [2, 4] (ContextWindowExceeded "context length exceeded") script + . runScriptLLMThrowingOn [2, 4] (ContextWindowExceeded "context length exceeded") script $ runProgram prog weatherQuestion res @?= Left (ContextWindowExceeded "context length exceeded") ] diff --git a/test/EnvSpec.hs b/test/EnvSpec.hs --- a/test/EnvSpec.hs +++ b/test/EnvSpec.hs @@ -7,8 +7,8 @@ import Data.Generics.Labels () import Data.List (find) import Data.Text qualified as T -import MockLLM (runEffMock) import Shikumi.Error (ShikumiError (..)) +import Shikumi.Testing (runEffScript) import Shikumi.Tool.Env ( DirEntry, ExecRequest (..), @@ -37,7 +37,7 @@ let file = T.pack (root "hello.txt") nested = T.pack (root "nested") result <- - runEffMock [] $ do + runEffScript [] $ do envWriteFile localToolEnv file "hello\n" bytes <- envReadFile localToolEnv file stat <- envStat localToolEnv file @@ -78,7 +78,7 @@ assertBool "cwd returns an absolute path" ("/" `T.isPrefixOf` cwd), testCase "a negative exec timeout is clamped, not disabled" $ do result <- - runEffMock [] $ + runEffScript [] $ envExec localToolEnv ExecRequest diff --git a/test/Fixtures.hs b/test/Fixtures.hs --- a/test/Fixtures.hs +++ b/test/Fixtures.hs @@ -33,10 +33,10 @@ import Data.Aeson (ToJSON, Value, object, (.=)) import Data.Text (Text) import GHC.Generics (Generic) -import MockLLM (mkTextResponse, mkToolCallResponse) import Shikumi.Adapter (ToPrompt) import Shikumi.Schema (FromModel, ToSchema, Validatable) import Shikumi.Signature (Signature, mkSignature) +import Shikumi.Testing (mkTextResponse, mkToolCallResponse) import Shikumi.Tool (SomeTool (..), Tool, ToolRegistry, mkRegistry, mkTool) -- --------------------------------------------------------------------------- diff --git a/test/FsSpec.hs b/test/FsSpec.hs --- a/test/FsSpec.hs +++ b/test/FsSpec.hs @@ -9,8 +9,8 @@ import Data.List (sort) import Data.Text (Text) import Data.Text qualified as T -import MockLLM (runEffMock) import Shikumi.Error (ShikumiError (..)) +import Shikumi.Testing (runEffScript) import Shikumi.Tool (Tool (..)) import Shikumi.Tool.Builtin.Fs ( EditReq (..), @@ -53,7 +53,7 @@ withTempDir "roundtrip" $ \root -> do let file = T.pack (root "note.txt") result <- - runEffMock [] $ do + runEffScript [] $ do rgPresent <- envExec localToolEnv (presentReq "rg") fdPresent <- envExec localToolEnv (presentReq "fd") written <- run (writeTool localToolEnv) WriteReq {path = file, content = "alpha\nold title\n"} @@ -94,7 +94,7 @@ nodeFile = T.pack (root "node_modules" "dep.txt") binaryFile = T.pack (root "binary.bin") result <- - runEffMock [] $ do + runEffScript [] $ do envWriteFile localToolEnv visible "SECRET visible\n" envMkdir localToolEnv (T.pack (root ".git")) envWriteFile localToolEnv gitFile "SECRET git\n" @@ -113,7 +113,7 @@ withTempDir "bad-regex" $ \root -> do let visible = T.pack (root "visible.txt") result <- - runEffMock [] $ do + runEffScript [] $ do envWriteFile localToolEnv visible "hello\n" run (grepTool noFastToolEnv) @@ -127,7 +127,7 @@ nestedDir = T.pack (root "sub") nested = T.pack (root "sub" "nested.txt") result <- - runEffMock [] $ do + runEffScript [] $ do envWriteFile localToolEnv top "marker top\n" envMkdir localToolEnv nestedDir envWriteFile localToolEnv nested "marker nested\n" @@ -155,7 +155,7 @@ withTempDir "read-truncated" $ \root -> do let file = T.pack (root "three.txt") result <- - runEffMock [] $ do + runEffScript [] $ do envWriteFile localToolEnv file "one\ntwo\nthree\n" offsetToEnd <- run (readTool localToolEnv) ReadReq {path = file, offset = Just 1, limit = Nothing} firstLine <- run (readTool localToolEnv) ReadReq {path = file, offset = Just 0, limit = Just 1} @@ -170,7 +170,7 @@ testCase "glob at exactly maxResults is not marked truncated" $ withTempDir "glob-cap" $ \root -> do result <- - runEffMock [] $ do + runEffScript [] $ do traverse_ ( \n -> envWriteFile @@ -193,7 +193,7 @@ Left _ -> assertBool "directory symlink creation unsupported; skipping" True Right () -> do result <- - runEffMock [] $ do + runEffScript [] $ do envWriteFile localToolEnv file "real\n" run (globTool noFastToolEnv) GlobReq {patternText = "*.txt", path = Just (T.pack root)} case result of diff --git a/test/Main.hs b/test/Main.hs --- a/test/Main.hs +++ b/test/Main.hs @@ -3,6 +3,7 @@ module Main (main) where import AcceptanceSpec qualified +import AgentHistorySpec qualified import BuiltinAcceptanceSpec qualified import CodeActSpec qualified import CompactionSpec qualified @@ -10,11 +11,16 @@ import FsSpec qualified import ProgramOfThoughtSpec qualified import ProtocolSpec qualified +import RLMSpec qualified import ReActSpec qualified +import ResponsesIntegrationSpec qualified +import ResponsesSchemaSpec qualified import RestrictedSpec qualified import SchemaSpec qualified +import SessionSpec qualified import ShellSpec qualified import Test.Tasty (defaultMain, testGroup) +import ToolOutputSpec qualified import ToolSpec qualified import WebSpec qualified @@ -23,8 +29,14 @@ defaultMain $ testGroup "shikumi-tools" - [ SchemaSpec.tests, + [ ResponsesSchemaSpec.tests, + ResponsesIntegrationSpec.tests, + SessionSpec.tests, + RLMSpec.tests, + AgentHistorySpec.tests, + SchemaSpec.tests, ToolSpec.tests, + ToolOutputSpec.tests, EnvSpec.tests, WebSpec.tests, FsSpec.tests, diff --git a/test/MockLLM.hs b/test/MockLLM.hs deleted file mode 100644 --- a/test/MockLLM.hs +++ /dev/null @@ -1,136 +0,0 @@ --- | A deterministic, network-free interpreter of the @LLM@ effect for the --- @shikumi-tools@ specs, plus builders for the two kinds of scripted assistant --- 'Response' a ReAct turn can produce: a plain-text reply (the prompt protocol and --- the extract step) and a native tool-call block (the native protocol). --- --- Each 'Shikumi.LLM.Complete' pops the next scripted 'Response' in order; an --- exhausted script yields an empty text response (specs always script enough). -module MockLLM - ( runMockLLM, - runMockLLMThrowingOnce, - runMockLLMThrowingOn, - runEffMock, - runAgent, - mkTextResponse, - mkUsageResponse, - mkToolCallResponse, - mkToolCallsResponse, - ) -where - -import Baikai - ( AssistantContent (..), - Model, - Response, - emptyResponse, - emptyTextContent, - emptyToolCall, - ) -import Control.Lens ((&), (.~)) -import Data.Aeson (Value) -import Data.Generics.Labels () -import Data.IORef (IORef, atomicModifyIORef', newIORef) -import Data.Text (Text) -import Data.Vector qualified as V -import Effectful (Eff, IOE, liftIO, runEff, type (:>)) -import Effectful.Dispatch.Dynamic (interpret) -import Effectful.Error.Static (Error, runErrorNoCallStack, throwError) -import Numeric.Natural (Natural) -import Shikumi.Error (ShikumiError) -import Shikumi.LLM (LLM (..)) -import Shikumi.Program (Program, runProgram) - --- | Interpret @LLM@ by popping scripted 'Response's in order. The streaming op is --- unused. An exhausted script yields an empty text response. -runMockLLM :: (IOE :> es) => [Response] -> Eff (LLM : es) a -> Eff es a -runMockLLM script act = do - ref <- liftIO (newIORef script) - interpret - ( \_ -> \case - Complete {} -> liftIO (pop ref) - Stream {} -> pure [] - ) - act - --- | Interpret @LLM@ like 'runMockLLM', but throw once on the first completion. -runMockLLMThrowingOnce :: - (IOE :> es, Error ShikumiError :> es) => - ShikumiError -> - [Response] -> - Eff (LLM : es) a -> - Eff es a -runMockLLMThrowingOnce = runMockLLMThrowingOn [1] - --- | Interpret @LLM@ like 'runMockLLM', but throw on selected 1-based completion --- calls. Useful for bounded-retry tests. -runMockLLMThrowingOn :: - (IOE :> es, Error ShikumiError :> es) => - [Int] -> - ShikumiError -> - [Response] -> - Eff (LLM : es) a -> - Eff es a -runMockLLMThrowingOn throwAt err script act = do - ref <- liftIO (newIORef script) - countRef <- liftIO (newIORef (0 :: Int)) - interpret - ( \_ -> \case - Complete {} -> do - n <- liftIO (atomicModifyIORef' countRef (\n0 -> let n1 = n0 + 1 in (n1, n1))) - if n `elem` throwAt - then throwError err - else liftIO (pop ref) - Stream {} -> pure [] - ) - act - --- | Discharge a network-free @LLM@ computation against a scripted mock: handle the --- error channel, then @IO@. The whole stack is exactly the row a ReAct agent runs --- in (@LLM@ + @Error ShikumiError@), plus @IOE@ at the bottom for the mock's @IORef@. -runEffMock :: - [Response] -> - Eff '[LLM, Error ShikumiError, IOE] a -> - IO (Either ShikumiError a) -runEffMock script = runEff . runErrorNoCallStack . runMockLLM script - --- | Run a 'Program' (e.g. a @react@ agent) against a scripted mock LM. -runAgent :: - [Response] -> - Program i o -> - i -> - IO (Either ShikumiError o) -runAgent script prog i = runEffMock script (runProgram prog i) - -pop :: IORef [Response] -> IO Response -pop ref = atomicModifyIORef' ref step - where - step (x : xs) = (xs, x) - step [] = ([], mkTextResponse "") - --- | An assistant 'Response' carrying @t@ as its single text block. -mkTextResponse :: Text -> Response -mkTextResponse t = - emptyResponse & #message . #content .~ V.singleton (AssistantText (emptyTextContent & #text .~ t)) - --- | A text response that also carries a resolved model and input-token usage. -mkUsageResponse :: Model -> Natural -> Text -> Response -mkUsageResponse model inputTokens text = - mkTextResponse text - & #model .~ model - & #message . #usage . #inputTokens .~ inputTokens - & #message . #usage . #totalTokens .~ inputTokens - --- | An assistant 'Response' carrying a single native tool-call block. -mkToolCallResponse :: Text -> Text -> Value -> Response -mkToolCallResponse callId nm args = mkToolCallsResponse [(callId, nm, args)] - --- | An assistant 'Response' carrying several native tool-call blocks in order. -mkToolCallsResponse :: [(Text, Text, Value)] -> Response -mkToolCallsResponse calls = - emptyResponse - & #message - . #content - .~ V.fromList - [ AssistantToolCall (emptyToolCall & #id_ .~ callId & #name .~ nm & #arguments .~ args) - | (callId, nm, args) <- calls - ] diff --git a/test/ProgramOfThoughtSpec.hs b/test/ProgramOfThoughtSpec.hs --- a/test/ProgramOfThoughtSpec.hs +++ b/test/ProgramOfThoughtSpec.hs @@ -12,7 +12,6 @@ import Data.Text (Text) import Data.Text qualified as T import GHC.Generics (Generic) -import MockLLM (mkTextResponse, runAgent) import Shikumi.Adapter (ToPrompt) import Shikumi.CodeExec.Interpreter (CodeInterpreter (..)) import Shikumi.CodeExec.ProgramOfThought (PoTConfig (..), programOfThought, programOfThoughtWith) @@ -21,6 +20,7 @@ import Shikumi.Program (foldParams) import Shikumi.Schema (FromModel, ToSchema, Validatable) import Shikumi.Signature (Signature, mkSignature) +import Shikumi.Testing (mkTextResponse, runAgent) import Test.Tasty (TestTree, testGroup) import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=)) diff --git a/test/ProtocolSpec.hs b/test/ProtocolSpec.hs --- a/test/ProtocolSpec.hs +++ b/test/ProtocolSpec.hs @@ -21,7 +21,6 @@ weatherRegistry, weatherSignature, ) -import MockLLM (mkTextResponse, mkToolCallsResponse, runAgent) import Shikumi.Agent.ReAct ( Action (..), Step (..), @@ -32,6 +31,7 @@ reactWithTrajectory, resolveProtocolKind, ) +import Shikumi.Testing (mkTextResponse, mkToolCallsResponse, runAgent) import Test.Tasty (TestTree, testGroup) import Test.Tasty.HUnit (assertFailure, testCase, (@?=)) diff --git a/test/RLMExample.hs b/test/RLMExample.hs new file mode 100644 --- /dev/null +++ b/test/RLMExample.hs @@ -0,0 +1,92 @@ +{-# LANGUAGE DeriveAnyClass #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE OverloadedLabels #-} +{-# LANGUAGE OverloadedStrings #-} + +module RLMExample (main, example) where + +import Baikai qualified as B +import Control.Lens ((&), (.~)) +import Data.Generics.Labels () +import Data.IORef (atomicModifyIORef', newIORef) +import Data.Text (Text) +import Data.Vector qualified as V +import Effectful (liftIO, runEff) +import Effectful.Dispatch.Dynamic (interpret) +import Effectful.Error.Static (runErrorNoCallStack, throwError) +import GHC.Generics (Generic) +import Shikumi.Adapter (ToPrompt) +import Shikumi.CodeExec.RLM +import Shikumi.CodeExec.Session +import Shikumi.Error (ShikumiError (..)) +import Shikumi.LLM (LLM (..)) +import Shikumi.Program (runProgram) +import Shikumi.Schema (FromModel, ToSchema, Validatable) +import Shikumi.Signature (Signature, mkSignature) + +newtype Question = Question {question :: Text} + deriving stock (Generic) + deriving anyclass (ToPrompt) + +newtype Answer = Answer {answer :: Text} + deriving stock (Generic, Show, Eq) + deriving anyclass (ToSchema, FromModel) + +instance Validatable Answer + +signature :: Signature Question Answer +signature = mkSignature "Inspect the source and answer the question." + +example :: IO (Either ShikumiError (Text, Int)) +example = do + replies <- + newIORef + [ "{\"op\":\"slice\",\"name\":\"source\",\"start\":0,\"count\":12}", + "{\"op\":\"query\",\"prompt\":\"What color is named in: Color: amber?\"}", + "amber", + "{\"op\":\"submit\",\"value\":{\"answer\":\"amber\"}}" + ] + runEff + . runErrorNoCallStack + . interpret + ( \_ -> \case + Complete {} -> do + next <- + liftIO + ( atomicModifyIORef' + replies + ( \case + [] -> ([], Nothing) + x : xs -> (xs, Just x) + ) + ) + case next of + Nothing -> throwError (ProviderFailure "Unexpected extra model call") + Just text -> pure (B.emptyResponse & #message . #content .~ V.singleton (B.AssistantText (B.emptyTextContent & #text .~ text))) + Stream {} -> throwError (ProviderFailure "Unexpected streaming call") + ) + $ do + -- Only the question is ToPrompt input. The source remains outside it. + source <- + either + (\(SessionError e) -> throwError (ValidationFailure e)) + pure + (contextStore [("source", "Color: amber")]) + (outcome, report) <- + runProgram + (rlmWithReport defaultRLMConfig source signature) + (Question "Which color is named?") + case outcome of + RLMSubmitted (Answer text) -> pure (text, subqueriesAttempted report) + RLMExhausted limit -> + throwError + ( BudgetExceeded + ( "Session exhausted: " <> case limit of + Subqueries -> "subqueries" + _ -> "another configured allowance" + ) + ) + +main :: IO () +main = example >>= print diff --git a/test/RLMSpec.hs b/test/RLMSpec.hs new file mode 100644 --- /dev/null +++ b/test/RLMSpec.hs @@ -0,0 +1,233 @@ +{-# LANGUAGE GADTs #-} + +module RLMSpec (tests) where + +import Baikai (Model, emptyModel) +import Baikai.Speed (Speed (..)) +import Control.Concurrent (forkFinally, newEmptyMVar, putMVar, takeMVar) +import Control.Lens ((&), (.~), (^.)) +import Data.Aeson (Value (..), object, (.=)) +import Data.Generics.Labels () +import Data.IORef (modifyIORef', newIORef, readIORef) +import Data.Text (Text) +import Data.Text qualified as T +import Effectful (liftIO, runEff) +import Effectful.Dispatch.Dynamic (interpret) +import Effectful.Error.Static (runErrorNoCallStack, throwError) +import GHC.Generics (Generic) +import RLMExample qualified +import Shikumi.Adapter (ToPrompt) +import Shikumi.CodeExec.Prompt (encodeText) +import Shikumi.CodeExec.RLM +import Shikumi.CodeExec.Session +import Shikumi.Error (ShikumiError (..)) +import Shikumi.LLM (LLM (..)) +import Shikumi.LLM.Defaults +import Shikumi.Program (Program, runProgram) +import Shikumi.Schema (FromModel, ToSchema, Validatable (..)) +import Shikumi.Signature (Signature, mkSignature) +import Shikumi.Testing (captureLLMRequests, mkTextResponse) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit + +newtype Question = Question {question :: Text} + deriving stock (Generic, Show) + deriving anyclass (ToPrompt) + +newtype Answer = Answer {answer :: Text} + deriving stock (Generic, Show, Eq) + deriving anyclass (ToSchema, FromModel) + +instance Validatable Answer where + validate (Answer a) = if T.null a then Left "answer must be nonempty" else Right (Answer a) + +sig :: Signature Question Answer +sig = mkSignature "Combine the two facts." + +-- Capture exact textual request content and selected models; fail if the loop +-- dispatches even one unscripted call (including hidden extraction or repair). +runCaptured :: [Either ShikumiError Text] -> Program Question a -> IO (Either ShikumiError a, [(Model, Text)]) +runCaptured script prog = do + effective <- newIORef [] + remaining <- newIORef script + captured <- newIORef [] + out <- + runEff + . runErrorNoCallStack @ShikumiError + . interpret + ( \_ -> \case + Complete model ctx _ -> do + -- Context JSON contains every message and system string, not just a + -- hand-selected excerpt. Decode-free capture also catches source leaks. + let rendered = encodeText (object ["system" .= (ctx ^. #systemPrompt), "messages" .= (ctx ^. #messages)]) + liftIO (modifyIORef' captured ((model, rendered) :)) + xs <- liftIO (readIORef remaining) + case xs of + [] -> throwError (ProviderFailure "unscripted call") + x : rest -> do + liftIO (modifyIORef' remaining (const rest)) + either throwError (pure . mkTextResponse) x + Stream {} -> throwError (ProviderFailure "unexpected stream") + ) + . captureLLMRequests effective + . withRequestDefaults (emptyRequestDefaults {defaultSpeed = Just SpeedFast}) + $ runProgram prog (Question "Find both facts") + requests <- readIORef effective + map (\(_, _, opts) -> opts ^. #speed) requests @?= replicate (length requests) (Just SpeedFast) + calls <- reverse <$> readIORef captured + pure (out, calls) + +docs :: [(Text, Text)] -> IO ContextStore +docs xs = either (\e -> assertFailure (show e) >> fail "docs") pure (contextStore xs) + +submit :: Text -> Text +submit a = encodeText (object ["op" .= ("submit" :: Text), "value" .= object ["answer" .= a]]) + +expectLimit :: SessionLimit -> Either ShikumiError (RLMOutcome Answer, RLMReport) -> Assertion +expectLimit limit (Right (out, report)) = do + out @?= RLMExhausted limit + termination report @?= Just limit +expectLimit _ x = assertFailure (show x) + +tests :: TestTree +tests = + testGroup + "experimental RLM" + [ testCase "documented RLM example" $ RLMExample.example >>= (@?= Right ("amber", 1)), + testCase "two-megabyte source, bounded requests, persisted facts and concurrent isolation" $ do + a <- newEmptyMVar + b <- newEmptyMVar + _ <- forkFinally (largeCase "amber" "owl") (putMVar a) + _ <- forkFinally (largeCase "violet" "fox") (putMVar b) + ra <- takeMVar a + rb <- takeMVar b + either (assertFailure . show) (@?= Right (Answer "amber owl")) ra + either (assertFailure . show) (@?= Right (Answer "violet fox")) rb, + testCase "three-question batch at two-call limit dispatches zero subqueries" $ do + d <- docs [] + let c = defaultRLMConfig {session = defaultSessionConfig {maxSubqueries = 2}} + (r, calls) <- runCaptured [Right "{\"op\":\"query_batch\",\"prompts\":[\"a\",\"b\",\"c\"]}"] (rlmWithReport c d sig) + expectLimit Subqueries r + length calls @?= 1 + case r of Right (_, rep) -> subqueriesAttempted rep @?= 0; _ -> pure (), + testCase "accepted batch executes exactly two ordered calls using selected model" $ do + d <- docs [] + let model = emptyModel & #modelId .~ "submodel" + c = defaultRLMConfig {session = defaultSessionConfig {maxSubqueries = 2, subModel = model}} + (r, calls) <- runCaptured (map Right ["{\"op\":\"query_batch\",\"prompts\":[\"first\",\"second\"]}", "one", "two", submit "one two"]) (rlmWithReport c d sig) + length calls @?= 4 + map ((^. #modelId) . fst) calls @?= ["", "submodel", "submodel", ""] + case r of + Right (RLMSubmitted (Answer "one two"), rep) -> do + subqueriesAttempted rep @?= 2 + operations rep @?= 2 + [observationValue o | row <- audit rep, Just o <- [auditObservation row]] @?= [String "one", String "two"] + _ -> assertFailure (show r), + testCase "failed subquery propagates infrastructure error and stops batch" $ do + d <- docs [] + (r, calls) <- runCaptured [Right "{\"op\":\"query_batch\",\"prompts\":[\"a\",\"b\"]}", Left (Timeout "cancelled")] (rlmWithReport defaultRLMConfig d sig) + r @?= Left (Timeout "cancelled") + length calls @?= 2, + testCase "malformed action and invalid submission recover without extraction" $ do + d <- docs [("empty", "")] + (r, calls) <- runCaptured (map Right ["{", "{\"op\":\"submit\",\"value\":{}}", submit "", submit "fixed"]) (rlmWithReport defaultRLMConfig d sig) + length calls @?= 4 + case r of + Right (RLMSubmitted (Answer "fixed"), rep) -> operations rep @?= 4 + _ -> assertFailure (show r), + testCase "outer, operation and observation exhaustion make no extra call" $ do + d <- docs [] + let bad = Right "{" + mapM_ + ( \(c, limit) -> do + (r, calls) <- runCaptured [bad] (rlmWithReport c d sig) + expectLimit limit r + length calls @?= 1 + ) + [ (defaultRLMConfig {maxOuterIterations = 1}, OuterIterations), + (defaultRLMConfig {session = defaultSessionConfig {maxOperations = 1}}, Operations), + (defaultRLMConfig {session = defaultSessionConfig {maxObservedChars = 1}}, ObservationCharacters) + ], + testCase "request bounds are enforced before outer or sub-model dispatch" $ do + d <- docs [] + (r, calls) <- runCaptured [] (rlmWithReport (defaultRLMConfig {maxRequestChars = 1}) d sig) + expectLimit RequestCharacters r + length calls @?= 0 + let c = defaultRLMConfig {maxRequestChars = 4000, session = defaultSessionConfig {maxSubqueryChars = 10000}} + query = encodeText (object ["op" .= ("query" :: Text), "prompt" .= T.replicate 5000 "q"]) + (r2, calls2) <- runCaptured [Right query] (rlmWithReport c d sig) + expectLimit RequestCharacters r2 + length calls2 @?= 1, + testCase "retained action history cannot exceed request ceiling" $ do + d <- docs [] + let c = defaultRLMConfig {maxRequestChars = 4000} + action = encodeText (object ["op" .= ("store" :: Text), "name" .= ("memo" :: Text), "value" .= T.replicate 3000 "a"]) + (r, calls) <- runCaptured [Right action] (rlmWithReport c d sig) + expectLimit RequestCharacters r + length calls @?= 1, + testCase "oversized subquery response is explicitly truncated in report and next request" $ do + d <- docs [] + let c = defaultRLMConfig {session = defaultSessionConfig {maxObservationChars = 200}} + (r, calls) <- runCaptured (map Right ["{\"op\":\"query\",\"prompt\":\"a\"}", T.replicate 10000 "😀", submit "ok"]) (rlmWithReport c d sig) + length calls @?= 3 + case r of + Right (RLMSubmitted _, rep) -> do + let obs = [o | row <- audit rep, Just o <- [auditObservation row]] + assertBool "truncation recorded" (any ((/= Nothing) . truncation) obs) + assertBool "all observations bounded" (all ((<= 200) . T.length . renderObservation) obs) + _ -> assertFailure (show r), + testCase "invalid configs fail before any call; convenience maps exhaustion" $ do + d <- docs [] + mapM_ + ( \c -> do + (r, calls) <- runCaptured [] (rlmWithReport c d sig) + case r of Left (ValidationFailure _) -> pure (); _ -> assertFailure (show r) + length calls @?= 0 + ) + [defaultRLMConfig {maxOuterIterations = 0}, defaultRLMConfig {maxRequestChars = (-1)}, defaultRLMConfig {session = defaultSessionConfig {maxSubqueries = 0}}] + (r, _) <- runCaptured [] (rlm (defaultRLMConfig {maxRequestChars = 1}) d sig) + r @?= Left (BudgetExceeded "RLM session exhausted: RequestCharacters") + ] + +largeCase :: Text -> Text -> IO (Either ShikumiError Answer) +largeCase first second = do + let prefix = "FACT_A=" <> first + source = prefix <> T.replicate (1500000 - T.length prefix) "." <> "FACT_B=" <> second <> T.replicate 600000 "." + c = defaultRLMConfig {maxRequestChars = 16000, session = defaultSessionConfig {maxScanChars = 2000000, maxObservationChars = 1000}} + action op fields = encodeText (object (("op" .= (op :: Text)) : fields)) + script = + map + Right + [ action "load" ["name" .= ("memo" :: Text)], + action "find" ["name" .= ("source" :: Text), "needle" .= ("FACT_A=" :: Text), "start" .= (0 :: Int), "limit" .= (1 :: Int)], + action "slice" ["name" .= ("source" :: Text), "start" .= (0 :: Int), "count" .= T.length prefix], + action "store" ["name" .= ("memo" :: Text), "value" .= first], + action "find" ["name" .= ("source" :: Text), "needle" .= ("FACT_B=" :: Text), "start" .= (1000000 :: Int), "limit" .= (1 :: Int)], + action "slice" ["name" .= ("source" :: Text), "start" .= (1500000 :: Int), "count" .= (7 + T.length second)], + action "load" ["name" .= ("memo" :: Text)], + action "query" ["prompt" .= ("Combine FACT_A=" <> first <> " and FACT_B=" <> second)], + first <> " " <> second, + submit (first <> " " <> second) + ] + d <- docs [("source", source)] + (r, calls) <- runCaptured script (rlmWithReport c d sig) + assertBool "fixture is at least two megabytes" (T.length source > 2000000) + length calls @?= 10 + mapM_ + ( \(_, request) -> do + assertBool "no source copied" (not (source `T.isInfixOf` request)) + assertBool "even serialized capture stays bounded" (T.length request <= maxRequestChars c) + ) + calls + case r of + Right (RLMSubmitted out, rep) -> do + operations rep @?= 9 + subqueriesAttempted rep @?= 1 + let obs = [o | row <- audit rep, Just o <- [auditObservation row]] + assertBool "fresh invocation has no memo" (any ((/= Nothing) . observationError) (take 1 obs)) + assertBool "first fact found at source origin" (any ((== object ["matches" .= ([0] :: [Int]), "scannedEnd" .= (2000000 :: Int), "nextOffset" .= (1 :: Int), "complete" .= False]) . observationValue) obs) + assertBool "second fact sliced at distant source offset" (any ((== object ["name" .= ("source" :: Text), "start" .= (1500000 :: Int), "end" .= (1500007 + T.length second), "text" .= ("FACT_B=" <> second)]) . observationValue) obs) + assertBool "stored fact is observed" (any ((== String first) . observationValue) obs) + pure (Right out) + Right other -> assertFailure (show other) >> pure (Left (ProviderFailure "exhausted")) + Left e -> pure (Left e) diff --git a/test/ReActSessionExample.hs b/test/ReActSessionExample.hs new file mode 100644 --- /dev/null +++ b/test/ReActSessionExample.hs @@ -0,0 +1,89 @@ +{-# LANGUAGE DeriveAnyClass #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE OverloadedLabels #-} +{-# LANGUAGE OverloadedStrings #-} + +-- | Offline example compiled and executed by the shikumi-tools test suite. +module ReActSessionExample (main, example) where + +import Baikai qualified as B +import Control.Lens ((&), (.~)) +import Data.Aeson (eitherDecode, encode, object, (.=)) +import Data.Generics.Labels () +import Data.IORef (atomicModifyIORef', newIORef) +import Data.Text (Text) +import Data.Text qualified as T +import Data.Vector qualified as V +import Effectful (liftIO, runEff) +import Effectful.Dispatch.Dynamic (interpret) +import Effectful.Error.Static (runErrorNoCallStack, throwError) +import GHC.Generics (Generic) +import Shikumi.Adapter (ToPrompt) +import Shikumi.Agent.History (decodeSession, encodeSession, sessionTurns) +import Shikumi.Agent.ReAct +import Shikumi.Error (ShikumiError (..)) +import Shikumi.LLM (LLM (..)) +import Shikumi.LLM.Continuation (validateRequestContinuation) +import Shikumi.Schema (FromModel, ToSchema, Validatable) +import Shikumi.Signature (Signature, mkSignature) +import Shikumi.Tool (mkDynTool, mkRegistry) +import Shikumi.Tool.Output (textToolOutput) + +newtype Question = Question {question :: Text} + deriving stock (Generic) + deriving anyclass (ToPrompt) + +newtype Answer = Answer {answer :: Text} + deriving stock (Generic, Show, Eq) + deriving anyclass (ToSchema, FromModel) + +instance Validatable Answer + +signature :: Signature Question Answer +signature = mkSignature "Use the lookup tool and answer the question." + +-- | Returns the typed answer and cumulative assistant exchange count. +example :: IO (Either ShikumiError (Text, Int)) +example = do + let registry = mkRegistry [mkDynTool "lookup" "Return the stored city." (object []) (\_ -> pure (Right (textToolOutput "Paris")))] + cfg = defaultReActConfig {protocol = ProtocolNative} + model = B.mkModel (B.Custom "session-example") "example" "https://example.invalid" + response callId name args = B.emptyResponse & #model .~ model & #message . #content .~ V.singleton (B.AssistantToolCall (B.ToolCall callId name args)) + script = [response "call-A" "lookup" (object []), response "final-A" finalToolName (object ["answer" .= ("Paris, France" :: Text)])] + replies <- newIORef script + runEff + . runErrorNoCallStack + . interpret + ( \_ -> \case + Complete m ctx opts -> do + either throwError pure (validateRequestContinuation m ctx opts) + next <- + liftIO + ( atomicModifyIORef' + replies + ( \xs -> case xs of + [] -> ([], Nothing) + x : rest -> (rest, Just x) + ) + ) + maybe (throwError (ProviderFailure "Unexpected extra model call")) pure next + Stream {} -> pure [] + ) + $ do + initial <- startSessionWithModel model signature registry cfg (Question "Which city is stored?") + first <- advanceSession signature registry cfg initial + checkpoint <- case first of + SessionPaused s -> pure s + _ -> throwError (ValidationFailure "Expected a checkpoint") + -- Write these bytes to caller-owned storage if persistence is needed. + let bytes = encode (encodeSession checkpoint) + decoded = eitherDecode bytes >>= either (Left . show) Right . decodeSession + saved <- either (throwError . ValidationFailure . T.pack) pure decoded + resumed <- continueSession signature registry cfg (Question "Include the country too.") saved + finished <- runSession signature registry cfg resumed + case finished of + SessionFinished (Answer text) s -> pure (text, sessionTurns s) + SessionPaused _ -> throwError (ValidationFailure "Iteration budget exhausted") + +main :: IO () +main = example >>= print diff --git a/test/ReActSpec.hs b/test/ReActSpec.hs --- a/test/ReActSpec.hs +++ b/test/ReActSpec.hs @@ -21,7 +21,6 @@ weatherRegistry, weatherSignature, ) -import MockLLM (mkTextResponse, runAgent) import Shikumi.Agent.ReAct ( Action (..), Step (..), @@ -31,6 +30,7 @@ reactWithTrajectory, ) import Shikumi.Error (ShikumiError (..)) +import Shikumi.Testing (mkTextResponse, runAgent) import Shikumi.Tool (SomeTool (..), Tool, ToolRegistry, mkRegistry, mkTool) import Test.Tasty (TestTree, testGroup) import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=)) diff --git a/test/ResponsesIntegrationSpec.hs b/test/ResponsesIntegrationSpec.hs new file mode 100644 --- /dev/null +++ b/test/ResponsesIntegrationSpec.hs @@ -0,0 +1,215 @@ +module ResponsesIntegrationSpec (tests) where + +import Baikai qualified as B +import Baikai.Error qualified as BE +import Control.Concurrent (forkFinally, killThread, threadDelay) +import Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar) +import Control.Lens ((&), (.~), (^.)) +import Data.Aeson +import Data.Aeson.KeyMap qualified as KM +import Data.Generics.Labels () +import Data.IORef +import Data.Text qualified as T +import Data.Vector qualified as V +import Effectful (Eff, IOE, liftIO, runEff, (:>)) +import Effectful.Concurrent (runConcurrent) +import Effectful.Dispatch.Dynamic (interpose) +import Effectful.Error.Static (runErrorNoCallStack, throwError) +import Effectful.Prim (runPrim) +import Shikumi.Agent.History qualified as H +import Shikumi.Agent.ReAct qualified as R +import Shikumi.Cache (cachedLLM) +import Shikumi.Cache.Backend.Memory (newMemoryCache, runCacheMemory) +import Shikumi.Compaction (CompactionConfig (..)) +import Shikumi.Effect.Time (runTime) +import Shikumi.Error (ShikumiError (..)) +import Shikumi.LLM qualified as L +import Shikumi.LLM.Defaults +import Shikumi.LLM.Observation qualified as O +import Shikumi.Routing (routeLLM, runRouting) +import Shikumi.Testing.Fixtures (Answer (..), Question (..), instructedSig) +import Shikumi.Testing.Responses +import Shikumi.Tool qualified as Tool +import Shikumi.Tool.Output (textToolOutput) +import System.Timeout (timeout) +import Test.Tasty +import Test.Tasty.HUnit + +cfg :: R.ReActConfig +cfg = R.defaultReActConfig {R.protocol = R.ProtocolNative, R.maxIters = 3, R.compaction = CompactionConfig 0 4 False} + +tools :: Tool.ToolRegistry +tools = + Tool.mkRegistry + [ Tool.mkDynTool "lookup" "Return the stored city." (object ["type" .= String "object", "properties" .= object []]) $ \_ -> do + _ <- L.complete B.emptyModel (B.emptyContext & #systemPrompt .~ Just "local-lookup") B.emptyOptions + pure (Right (textToolOutput "Paris")) + ] + +-- Count actual body entries via a local effect handler; it never dispatches HTTP. +countLookup :: (L.LLM :> es, IOE :> es) => IORef Int -> Eff es a -> Eff es a +countLookup ref = interpose $ \_ -> \case + L.Complete _ c _ | c ^. #systemPrompt == Just "local-lookup" -> do + liftIO (atomicModifyIORef' ref (\n -> (n + 1, ()))) + pure B.emptyResponse + L.Complete m c o -> L.complete m c o + L.Stream m c o -> L.stream m c o + +firstReply :: Reply +firstReply = sseReply [completed [reasoningItem, functionItem "item-lookup" "call-lookup" "lookup" (object [])]] + +lastReply :: Reply +lastReply = sseReply [completed [functionItem "item-final" "call-final" R.finalToolName (object ["answer" .= String "Paris", "confidence" .= (0.9 :: Double)])]] + +restore :: H.ReActSession -> Either String H.ReActSession +restore s = eitherDecode (encode (H.encodeSession s)) >>= either (Left . show) Right . H.decodeSession + +tests :: TestTree +tests = + testGroup + "Responses sessions" + [ testCase "real reasoning and tool IDs survive checkpoint bytes with one dispatch" $ + withResponsesFixture [firstReply, lastReply] $ \f -> do + ref <- newIORef 0 + result <- runEff + . runErrorNoCallStack @ShikumiError + . runRouting (model f) + . L.runLLMWith (registry f) + . withTransportOptions fixtureOptions + . withRequestDefaults (emptyRequestDefaults {defaultThinking = Just B.ThinkingHigh}) + . routeLLM + . countLookup ref + $ do + initial <- R.startSessionWithModel (model f) instructedSig tools cfg (Question "Look up the city.") + step <- R.advanceSession instructedSig tools cfg initial + checkpoint <- case step of R.SessionPaused s -> pure s; _ -> throwError (ValidationFailure "expected checkpoint") + saved <- either (throwError . ValidationFailure . T.pack) pure (restore checkpoint) + resumed <- R.continueSession instructedSig tools cfg (Question "Submit the answer.") saved + finished <- R.runSession instructedSig tools cfg resumed + pure (checkpoint, saved, finished) + case result of + Right (before, saved, R.SessionFinished answer _) -> do + before @?= saved + answer @?= Answer "Paris" 0.9 + assertBool "unsafe compaction fails" (either (const True) (const False) (H.compactSession 1 "summary" saved)) + other -> assertFailure (show other) + readIORef ref >>= (@?= 1) + bodies <- requests f + length bodies @?= 2 + case bodies of + [_, Object second] -> case KM.lookup "input" second of + Just (Array items) -> do + assertBool "opaque reasoning item is byte-structurally intact" (reasoningItem `elem` V.toList items) + let typed kind = [o | Object o <- V.toList items, KM.lookup "type" o == Just (String kind)] + map (KM.lookup "call_id") (typed "function_call") @?= [Just (String "call-lookup")] + map (KM.lookup "call_id") (typed "function_call_output") @?= [Just (String "call-lookup")] + let kinds = [KM.lookup "type" o | Object o <- V.toList items] + assertBool "call precedes result" (V.findIndex (== Just (String "function_call")) (V.fromList kinds) < V.findIndex (== Just (String "function_call_output")) (V.fromList kinds)) + _ -> assertFailure "missing wire input" + _ -> assertFailure "wrong request count", + testCase "wrong routed model and API fail before a second HTTP request" $ + withResponsesFixture [firstReply] $ \f -> do + ref <- newIORef 0 + initial <- runEff + . runErrorNoCallStack @ShikumiError + . runRouting (model f) + . L.runLLMWith (registry f) + . withTransportOptions fixtureOptions + . routeLLM + . countLookup ref + $ do + s <- R.startSessionWithModel (model f) instructedSig tools cfg (Question "lookup") + R.advanceSession instructedSig tools cfg s + case initial of + Right (R.SessionPaused checkpoint) -> + mapM_ + ( \wrong -> do + result <- + runEff + . runErrorNoCallStack @ShikumiError + . runRouting wrong + . L.runLLMWith (registry f) + . withTransportOptions fixtureOptions + . routeLLM + $ R.advanceSession instructedSig tools cfg checkpoint + assertBool "origin rejected" (case result of Left (ValidationFailure _) -> True; _ -> False) + ) + [model f & #modelId .~ "wrong-model", model f & #api .~ B.OpenAIChatCompletions] + other -> assertFailure (show other) + requests f >>= (\xs -> length xs @?= 1) + readIORef ref >>= (@?= 1), + testCase "released refusal error is terminal; transient transport retries are observed" $ do + withResponsesFixture [sseReply [failureFrame "content_filter"]] $ \f -> do + (observe, snapshot) <- O.newBillingCollectorWithLimit 4 + let policy = (L.defaultLLMConfig (registry f)) {L.observer = Just observe, L.retryPolicy = L.RetryPolicy 3 0 0} + result <- + runEff . runConcurrent . runErrorNoCallStack @ShikumiError . L.runLLMResilient policy $ + L.complete (model f) B.emptyContext fixtureOptions + case result of + Left (ProviderError e) -> BE.category e @?= BE.OtherError + other -> assertFailure (show other) + requests f >>= (\xs -> length xs @?= 1) + snapshot >>= (\s -> O.failedAttempts s @?= 1) + withResponsesFixture [Reply 500 [] 0, sseReply [completed [messageItem "ok"]]] $ \f -> do + (observe, snapshot) <- O.newBillingCollectorWithLimit 4 + let policy = (L.defaultLLMConfig (registry f)) {L.observer = Just observe, L.retryPolicy = L.RetryPolicy 2 0 0} + result <- + runEff . runConcurrent . runErrorNoCallStack @ShikumiError . L.runLLMResilient policy $ + L.complete (model f) B.emptyContext fixtureOptions + assertBool "retry succeeds" (either (const False) (const True) result) + summary <- snapshot + map O.attempt (O.retainedAttempts summary) @?= [1, 2] + O.completedAttempts summary @?= 1 + O.failedAttempts summary @?= 1, + testCase "cache hit adds no released transport attempt" $ + withResponsesFixture [sseReply [completed [messageItem "ok"]]] $ \f -> do + cache <- newMemoryCache + (observe, snapshot) <- O.newBillingCollectorWithLimit 4 + let policy = (L.defaultLLMConfig (registry f)) {L.observer = Just observe} + result <- runEff + . runPrim + . runTime + . runConcurrent + . runErrorNoCallStack @ShikumiError + . runCacheMemory cache + . L.runLLMResilient policy + . cachedLLM + $ do + _ <- L.complete (model f) B.emptyContext fixtureOptions + _ <- L.complete (model f) B.emptyContext fixtureOptions + pure () + result @?= Right () + requests f >>= (\xs -> length xs @?= 1) + snapshot >>= (\s -> O.completedAttempts s @?= 1), + testCase "timeout and cancellation release loopback workers" $ do + cleanup <- withResponsesFixture [Reply 200 [completed [messageItem "late"]] 2000000] $ \f -> do + result <- + runEff . runErrorNoCallStack @ShikumiError . L.runLLMWith (registry f) $ + L.complete (model f) B.emptyContext (fixtureOptions & #timeoutMs .~ Just 50) + assertBool "transport timeout fails" (either (const True) (const False) result) + pure (activeRequests f) + drained <- timeout 2000000 (waitForZero cleanup) + drained @?= Just () + cleanup2 <- withResponsesFixture [Reply 200 [completed [messageItem "late"]] 2000000] $ \f -> do + done <- newEmptyMVar + tid <- + forkFinally + ( runEff . runErrorNoCallStack @ShikumiError . L.runLLMWith (registry f) $ + L.complete (model f) B.emptyContext (fixtureOptions & #timeoutMs .~ Just 5000) + ) + (putMVar done) + entered <- timeout 2000000 (waitForActive (activeRequests f)) + killThread tid + outcome <- takeMVar done + assertBool "cancellation propagates from the released adapter" (either (const True) (const False) outcome) + entered @?= Just () + pure (activeRequests f) + drained2 <- timeout 2000000 (waitForZero cleanup2) + drained2 @?= Just () + ] + +waitForZero :: IO Int -> IO () +waitForZero readCount = readCount >>= \n -> if n == 0 then pure () else threadDelay 1000 >> waitForZero readCount + +waitForActive :: IO Int -> IO () +waitForActive readCount = readCount >>= \n -> if n > 0 then pure () else threadDelay 1000 >> waitForActive readCount diff --git a/test/ResponsesSchemaSpec.hs b/test/ResponsesSchemaSpec.hs new file mode 100644 --- /dev/null +++ b/test/ResponsesSchemaSpec.hs @@ -0,0 +1,91 @@ +module ResponsesSchemaSpec (tests) where + +import Baikai qualified as B +import Control.Lens ((&), (.~), (^.)) +import Data.Aeson (Value (..), object, (.=)) +import Data.Aeson.Key (Key) +import Data.Aeson.KeyMap qualified as KM +import Data.Generics.Labels () +import Data.Vector qualified as V +import Effectful (runEff) +import Effectful.Error.Static (runErrorNoCallStack) +import Shikumi.Adapter (attachSchema) +import Shikumi.Error (ShikumiError) +import Shikumi.LLM qualified as L +import Shikumi.LLM.Defaults +import Shikumi.Program (runProgram) +import Shikumi.Routing (routeLLM, runRouting) +import Shikumi.Schema (deriveSchema) +import Shikumi.Testing.Fixtures +import Shikumi.Testing.Responses +import Test.Tasty +import Test.Tasty.HUnit + +field :: [Key] -> Value -> Maybe Value +field [] v = Just v +field (k : ks) (Object o) = KM.lookup k o >>= field ks +field _ _ = Nothing + +tests :: TestTree +tests = + testGroup + "Responses real adapter" + [ testCase "typed schema, defaults and streamed terminal cross actual HTTP" $ + withResponsesFixture (replicate 2 (sseReply [completed [messageItem "{\"answer\":\"Forty-two\",\"confidence\":0.9}"]])) $ \f -> do + let defaults = emptyRequestDefaults {defaultThinking = Just B.ThinkingHigh, defaultMaxTokens = Just 128} + schema = object ["type" .= String "object", "properties" .= object ["answer" .= object ["type" .= String "string"]]] + result <- + runEff + . runErrorNoCallStack @ShikumiError + . runRouting (model f) + . L.runLLMWith (registry f) + . withTransportOptions fixtureOptions + . withRequestDefaults defaults + . routeLLM + $ runProgram instructedProg (Question "What is the answer?") + result @?= Right (Answer "Forty-two" 0.9) + streamed <- + runEff + . runErrorNoCallStack @ShikumiError + . runRouting (model f) + . L.runLLMWith (registry f) + . withTransportOptions fixtureOptions + . withRequestDefaults defaults + . routeLLM + $ L.stream (model f) B.emptyContext (attachSchema schema B.emptyOptions) + case streamed of + Right events -> case [p | B.EventDone tp <- events, B.AssistantMessage p <- [tp ^. #message]] of + [p] -> do + p ^. #usage . #inputTokens @?= 20 + p ^. #usage . #outputTokens @?= 5 + _ -> assertFailure "missing terminal payload" + Left e -> assertFailure (show e) + bodies <- requests f + length bodies @?= 2 + map (field ["text", "format", "schema"]) bodies @?= [Just (deriveSchema @Answer), Just schema] + mapM_ + ( \body -> do + field ["model"] body @?= Just (String "fixture-reasoner") + field ["stream"] body @?= Just (Bool True) + field ["store"] body @?= Just (Bool False) + field ["max_output_tokens"] body @?= Just (Number 128) + field ["reasoning", "effort"] body @?= Just (String "high") + field ["text", "format", "type"] body @?= Just (String "json_schema") + field ["text", "format", "strict"] body @?= Just (Bool True) + field ["include"] body @?= Just (Array (V.singleton (String "reasoning.encrypted_content"))) + ) + bodies, + testCase "unsupported stop and image tool output fail before HTTP" $ + withResponsesFixture [] $ \f -> do + let badImage = B.ToolResultMessage (B.ToolResultPayload "call" "lookup" (V.singleton (B.ToolResultImage (B.ImageContent "bytes" "image/png"))) False Nothing) + -- Both failures are exercised through the released mapper. + a <- + runEff . runErrorNoCallStack @ShikumiError . L.runLLMWith (registry f) $ + L.complete (model f) B.emptyContext (fixtureOptions & #stopSequences .~ ["stop"]) + b <- + runEff . runErrorNoCallStack @ShikumiError . L.runLLMWith (registry f) $ + L.complete (model f) (B.emptyContext & #messages .~ V.singleton badImage) fixtureOptions + assertBool "stop rejected" (either (const True) (const False) a) + assertBool "image result rejected" (either (const True) (const False) b) + requests f >>= (@?= []) + ] diff --git a/test/SessionSpec.hs b/test/SessionSpec.hs new file mode 100644 --- /dev/null +++ b/test/SessionSpec.hs @@ -0,0 +1,113 @@ +module SessionSpec (tests) where + +import Data.Aeson (Value (..), object, (.=)) +import Data.Text qualified as T +import Shikumi.CodeExec.Session +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit + +setup :: SessionConfig -> [(T.Text, T.Text)] -> IO SessionState +setup c ds = either (\e -> assertFailure (show e) >> fail "setup") pure (contextStore ds >>= newSession c) + +observation :: Either SessionLimit SessionResult -> IO SessionObservation +observation (Right (Observed o)) = pure o +observation x = assertFailure (show x) >> fail "observation" + +tests :: TestTree +tests = + testGroup + "bounded session" + [ testCase "variables persist, sources cannot be shadowed, sessions isolate" $ do + s <- setup defaultSessionConfig [("doc", "hello")] + let (s1, _) = stepSession s (Right (Store "memo" (String "remember"))) + (s2, rejected) = stepSession s1 (Right (Store "doc" Null)) + o <- observation rejected + assertBool "shadow rejected" (observationError o /= Nothing) + snd (stepSession s2 (Right (Load "memo"))) @?= Right (Observed (SessionObservation (String "remember") Nothing Nothing)) + isolated <- observation (snd (stepSession s (Right (Load "memo")))) + assertBool "private variable" (observationError isolated /= Nothing) + operationCount s2 @?= 2, + testCase "Unicode slices use character offsets and empty documents work" $ do + s <- setup defaultSessionConfig [("doc", "a😀界z"), ("empty", "")] + o <- observation (snd (stepSession s (Right (Slice "doc" 1 2)))) + observationValue o @?= object ["name" .= ("doc" :: T.Text), "start" .= (1 :: Int), "end" .= (3 :: Int), "text" .= ("😀界" :: T.Text)] + e <- observation (snd (stepSession s (Right (Slice "empty" 0 0)))) + observationError e @?= Nothing, + testCase "invalid offsets, unknown names and builtins consume operations" $ do + s <- setup defaultSessionConfig [("doc", "abc")] + mapM_ + ( \a -> do + let (s1, r) = stepSession s (Right a) + o <- observation r + assertBool (show a) (observationError o /= Nothing) + operationCount s1 @?= 1 + ) + [Slice "doc" (-1) 1, Slice "doc" 4 0, Slice "doc" 1 3, Find "doc" "" 0 1, Find "doc" "a" 0 0, Load "absent", Describe "absent", Store "slice" Null], + testCase "find scans bounded text and supports overlap and continuation" $ do + s <- setup (defaultSessionConfig {maxScanChars = 4}) [("doc", "aaaaXYZ")] + o <- observation (snd (stepSession s (Right (Find "doc" "aa" 0 2)))) + observationValue o @?= object ["matches" .= ([0, 1] :: [Int]), "scannedEnd" .= (4 :: Int), "nextOffset" .= (2 :: Int), "complete" .= False], + testCase "stored character budget includes all values and names; rejection is atomic" $ do + s <- setup (defaultSessionConfig {maxStoredChars = 10}) [] + let (s1, _) = stepSession s (Right (Store "x" (String "ok"))) + (s2, r) = stepSession s1 (Right (Store "x" (String (T.replicate 20 "x")))) + o <- observation r + assertBool "oversized rejected" (observationError o /= Nothing) + snd (stepSession s2 (Right (Load "x"))) @?= snd (stepSession s1 (Right (Load "x"))), + testCase "all limits must be positive and oversized contexts are refused" $ do + let updates = [\c -> c {maxContextChars = 0}, \c -> c {maxStoredChars = (-1)}, \c -> c {maxActionBytes = 0}, \c -> c {maxObservationChars = 0}, \c -> c {maxObservedChars = 0}, \c -> c {maxOperations = 0}, \c -> c {maxScanChars = 0}, \c -> c {maxMatches = 0}, \c -> c {maxSubqueries = 0}, \c -> c {maxSubqueryChars = 0}] + mapM_ (\f -> assertBool "invalid config" (validateSessionConfig (f defaultSessionConfig) /= Right ())) updates + case contextStore [("d", "oversized")] >>= newSession (defaultSessionConfig {maxContextChars = 2}) of + Left _ -> pure () + Right _ -> assertFailure "context admitted", + testCase "UTF8 byte limit and malformed JSON are recoverable" $ do + let c = defaultSessionConfig {maxActionBytes = 10} + s <- setup c [] + mapM_ + ( \raw -> do + let (s1, r) = stepSession s (parseSessionAction c raw) + o <- observation r + assertBool "parse error" (observationError o /= Nothing) + operationCount s1 @?= 1 + ) + ["😀😀😀", "{", "{\"op\":1}"], + testCase "observation bounds include escaping and truncation metadata" $ do + s <- setup (defaultSessionConfig {maxObservationChars = 200, maxObservedChars = 250}) [] + let (s1, r) = recordObservation s (SessionObservation (String (T.replicate 1000 "\n")) Nothing Nothing) + case r of + Left e -> assertFailure (show e) + Right o -> do + assertBool "marked" (truncation o /= Nothing) + assertBool "serialized bound" (T.length (renderObservation o) <= 200) + observedChars s1 @?= T.length (renderObservation o) + snd (recordObservation s1 o) @?= Left ObservationCharacters, + testCase "truncated Unicode slices retain source offsets and can resume" $ do + s <- setup (defaultSessionConfig {maxObservationChars = 250}) [("doc", T.replicate 1000 "😀\n")] + o <- observation (snd (stepSession s (Right (Slice "doc" 0 1000)))) + case snd (recordObservation s o) of + Right bounded -> case truncation bounded of + Just tr -> do + assertBool "prefix is nonempty" (displayedChars tr > 0) + nextOffset tr @?= Just (displayedChars tr) + assertBool "serialized slice is bounded" (T.length (renderObservation bounded) <= 250) + next <- observation (snd (stepSession s (Right (Slice "doc" (displayedChars tr) 2)))) + observationError next @?= Nothing + Nothing -> assertFailure "missing slice truncation" + Left e -> assertFailure (show e), + testCase "subquery slots persist across actions and oversized prompts reserve none" $ do + s <- setup (defaultSessionConfig {maxSubqueries = 2, maxSubqueryChars = 3}) [] + let (s1, r) = stepSession s (Right (Query "long")) + o <- observation r + assertBool "prompt rejected" (observationError o /= Nothing) + let (s2, r2) = stepSession s1 (Right (Query "a")) + r2 @?= Right (RunQueries ["a"]) + snd (stepSession s2 (Right (QueryBatch ["b", "c"]))) @?= Left Subqueries, + testCase "operation exhaustion and whole batch reservation are exact" $ do + s <- setup (defaultSessionConfig {maxOperations = 1, maxSubqueries = 2}) [] + snd (stepSession s (Right (QueryBatch ["a", "b", "c"]))) @?= Left Subqueries + let (s1, r) = stepSession s (Right (QueryBatch ["a", "b"])) + r @?= Right (RunQueries ["a", "b"]) + subqueryAttempts s1 @?= 0 + subqueryAttempts (attemptedSubquery (attemptedSubquery s1)) @?= 2 + snd (stepSession s1 (Right (Submit Null))) @?= Left Operations + ] diff --git a/test/ShellSpec.hs b/test/ShellSpec.hs --- a/test/ShellSpec.hs +++ b/test/ShellSpec.hs @@ -4,7 +4,7 @@ import Control.Lens ((^.)) import Data.Text qualified as T -import MockLLM (runEffMock) +import Shikumi.Testing (runEffScript) import Shikumi.Tool (Tool (..)) import Shikumi.Tool.Builtin.Shell (BashReq (..), bashTool) import Shikumi.Tool.Env (localToolEnv) @@ -17,7 +17,7 @@ "Tool.Shell" [ testCase "bash captures stdout and zero exit" $ do result <- - runEffMock [] $ + runEffScript [] $ run (bashTool localToolEnv) BashReq {command = "echo hi", cwd = Nothing, timeoutMs = Just 5000, stdin = Nothing} @@ -28,7 +28,7 @@ assertBool "stdout contains hi" ("hi" `T.isInfixOf` (resp ^. #stdout)), testCase "bash returns stderr and non-zero exit as a value" $ do result <- - runEffMock [] $ + runEffScript [] $ run (bashTool localToolEnv) BashReq {command = "echo oops 1>&2; exit 3", cwd = Nothing, timeoutMs = Just 5000, stdin = Nothing} diff --git a/test/ToolOutputSpec.hs b/test/ToolOutputSpec.hs new file mode 100644 --- /dev/null +++ b/test/ToolOutputSpec.hs @@ -0,0 +1,55 @@ +module ToolOutputSpec (tests) where + +import Baikai qualified as B +import Control.Lens ((&), (.~), (^.)) +import Data.Aeson (object, (.=)) +import Data.Generics.Labels () +import Data.Text qualified as T +import Data.Vector qualified as V +import Effectful.Error.Static (throwError) +import Fixtures (weatherArgs, weatherRegistry) +import Shikumi.Error (ShikumiError (..)) +import Shikumi.Testing (runEffScript) +import Shikumi.Tool +import Shikumi.Tool.Output +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit + +tests :: TestTree +tests = + testGroup + "ToolOutput" + [ testCase "typed text compatibility" $ do + let call = B.ToolCall "A" "get_weather" weatherArgs + rich <- runEffScript [] (runToolCallOutput weatherRegistry call) + plain <- runEffScript [] (runToolCall weatherRegistry call) + fmap (fmap renderToolOutput) rich @?= plain, + testCase "structured-only dynamic output is preserved and visible" $ do + let value = object ["answer" .= (42 :: Int)] + output = ToolOutput (B.ToolResult V.empty False) (Just value) [value] + reg = mkRegistry [mkDynTool "dynamic" "runtime" (object []) (\_ -> pure (Right output))] + call = B.ToolCall "A" "dynamic" (object []) + actual <- runEffScript [] (runToolCallOutput reg call) + actual @?= Right (Right output) + assertBool "JSON labels" ("Structured JSON:" `T.isInfixOf` renderToolOutput output && "Extension JSON:" `T.isInfixOf` renderToolOutput output) + case toolOutputMessage call output of + B.ToolResultMessage p -> do + p ^. #toolCallId @?= "A" + V.length (p ^. #content) @?= 2 + _ -> assertFailure "expected tool message", + testCase "mixed native blocks retain images and error flag" $ do + let img = B.emptyImageContent & #imageData .~ "bytes" & #mimeType .~ "image/png" + blocks = V.fromList [B.ToolResultText (B.TextContent "caption"), B.ToolResultImage img] + output = ToolOutput (B.ToolResult blocks True) Nothing [] + case toolOutputMessage (B.ToolCall "A" "tool" (object [])) output of + B.ToolResultMessage p -> do + p ^. #content @?= blocks + p ^. #isError @?= True + _ -> assertFailure "expected tool message" + assertBool "error flag visible in prompt/text projection" ("Tool error:" `T.isPrefixOf` renderToolOutput output) + assertBool "image not dropped in text projection" ("Image JSON:" `T.isInfixOf` renderToolOutput output), + testCase "dynamic recoverable failure is model-visible" $ do + let tool = mkDynTool "bad" "" (object []) (\_ -> throwError (ValidationFailure "bad result")) + result <- runEffScript [] (runErasedOutput tool (object [])) + result @?= Right (Left (ToolRunFailed "bad" "bad result")) + ] diff --git a/test/ToolSpec.hs b/test/ToolSpec.hs --- a/test/ToolSpec.hs +++ b/test/ToolSpec.hs @@ -13,8 +13,8 @@ import Data.Text qualified as T import Effectful.Error.Static (throwError) import Fixtures (WeatherReq, WeatherResp, weatherArgs, weatherRegistry) -import MockLLM (runEffMock) import Shikumi.Error (ShikumiError (..)) +import Shikumi.Testing (runEffScript) import Shikumi.Tool (SomeTool (..), Tool, ToolError (..), mkRegistry, mkTool, runToolCall) import Test.Tasty (TestTree, testGroup) import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=)) @@ -37,27 +37,27 @@ testGroup "Tool" [ testCase "decodes valid args and runs the body" $ do - res <- runEffMock [] (runToolCall weatherRegistry (tc "get_weather" weatherArgs)) + res <- runEffScript [] (runToolCall weatherRegistry (tc "get_weather" weatherArgs)) case res of Right (Right obs) -> assertBool "observation mentions the forecast" ("mild" `T.isInfixOf` obs) other -> assertFailure ("expected a tool result, got " <> show other), testCase "returns ToolArgsInvalid for a missing required field" $ do - res <- runEffMock [] (runToolCall weatherRegistry (tc "get_weather" (object ["city" .= ("Paris" :: Text)]))) + res <- runEffScript [] (runToolCall weatherRegistry (tc "get_weather" (object ["city" .= ("Paris" :: Text)]))) case res of Right (Left (ToolArgsInvalid nm _)) -> nm @?= "get_weather" other -> assertFailure ("expected ToolArgsInvalid, got " <> show other), testCase "returns ToolNotFound for an unknown name" $ do - res <- runEffMock [] (runToolCall weatherRegistry (tc "nope" (object []))) + res <- runEffScript [] (runToolCall weatherRegistry (tc "nope" (object []))) case res of Right (Left (ToolNotFound nm)) -> nm @?= "nope" other -> assertFailure ("expected ToolNotFound, got " <> show other), testCase "a tool body throwing BudgetExceeded escapes as ShikumiError" $ do - res <- runEffMock [] (runToolCall (mkRegistry [SomeTool budgetTool]) (tc "burn_budget" weatherArgs)) + res <- runEffScript [] (runToolCall (mkRegistry [SomeTool budgetTool]) (tc "burn_budget" weatherArgs)) case res of Left (BudgetExceeded msg) -> msg @?= "ceiling reached" other -> assertFailure ("expected escaped BudgetExceeded, got " <> show other), testCase "a tool body throwing ValidationFailure becomes ToolRunFailed" $ do - res <- runEffMock [] (runToolCall (mkRegistry [SomeTool flakyTool]) (tc "flaky" weatherArgs)) + res <- runEffScript [] (runToolCall (mkRegistry [SomeTool flakyTool]) (tc "flaky" weatherArgs)) case res of Right (Left (ToolRunFailed nm msg)) -> do nm @?= "flaky" diff --git a/test/WebSpec.hs b/test/WebSpec.hs --- a/test/WebSpec.hs +++ b/test/WebSpec.hs @@ -11,8 +11,8 @@ import Data.IORef (IORef, atomicModifyIORef', newIORef, readIORef) import Data.Text (Text) import Data.Text qualified as T -import MockLLM (runEffMock) import Shikumi.Error (ShikumiError (..)) +import Shikumi.Testing (runEffScript) import Shikumi.Tool (SomeTool (..), ToolRegistry, mkRegistry, runToolCall) import Shikumi.Tool.Builtin.Web (webFetchTool, webSearchTool) import Shikumi.Tool.Web @@ -39,7 +39,7 @@ "Tool.Web" [ testCase "web_fetch returns a stubbed 200 observation" $ do result <- - runEffMock [] $ + runEffScript [] $ runToolCall stubRegistry (tc "web_fetch" (object ["url" .= ("https://example.test" :: Text)])) @@ -50,7 +50,7 @@ other -> assertFailure ("expected web_fetch observation, got " <> show other), testCase "web_fetch surfaces a 404 status as a value" $ do result <- - runEffMock [] $ + runEffScript [] $ runToolCall notFoundRegistry (tc "web_fetch" (object ["url" .= ("https://example.test/missing" :: Text)])) @@ -61,7 +61,7 @@ other -> assertFailure ("expected 404 as a tool value, got " <> show other), testCase "web_search returns stubbed hits" $ do result <- - runEffMock [] $ + runEffScript [] $ runToolCall stubRegistry (tc "web_search" (object ["query" .= ("shikumi" :: Text)])) @@ -93,7 +93,7 @@ testCase "fetch of a denied URL fails fast with ValidationFailure" $ do manager <- newTlsManager result <- - runEffMock [] $ + runEffScript [] $ webFetch (localWebClient manager Nothing) "http://169.254.169.254/latest/meta-data/" @@ -109,7 +109,7 @@ Just _ -> do manager <- newTlsManager result <- - runEffMock [] $ + runEffScript [] $ runToolCall (mkRegistry [SomeTool (webFetchTool (localWebClient manager Nothing))]) (tc "web_fetch" (object ["url" .= ("https://example.com" :: Text)]))