Coverage for core/langgraph/_step_functions.py: 99%

1""" 

2_step_functions.py — StepFunction registry consultation and execution (ADR-0037). 

3 

4`kind: StepFunction` is the deterministic, consumer-authored code step. It is 

5invoked through the same step form as a platform builtin (`type: builtin`, 

6`component: <name>`); the engine resolves the component name against the 

7builtin register first and then against the StepFunction registry (ADR-0037 

8D2 — a name present in both is a hard error, enforced at the dispatch site). 

9 

10Runtime contract (ADR-0037 D4): the declared `entrypoint` runs with the 

11declared `runtime` (`bash` | `python`); the step's resolved `params` are fed 

12as JSON on stdin; stdout must be a JSON object whose keys are the declared 

13`spec.interface.outputs` — required outputs must be present, optional outputs 

14may be absent, and undeclared keys are rejected. The declared outputs become 

15the flat step result `state[step_id][<key>]`. A non-zero exit code, non-JSON 

16stdout, or output-contract violation is a hard run failure (ADR-0012). 

17 

18Transport: the entrypoint always reads params JSON on stdin. Locally the 

19engine feeds it by direct stdin; in-cluster the params are written to a 

20shared-workspace file and the sidecar command redirects that file into stdin 

21(`< <params_file>`), because the K8s exec v4 client has no half-close. 

22 

23The class is constructed over the existing kind-generic `ManifestRegistry` 

24and passed to `StepRunner` by `GraphBuilder`. 

25""" 

26from __future__ import annotations 

27 

28import contextlib 

29import json 

30import logging 

31import os 

32import re 

33import shlex 

34import sys 

35from pathlib import Path 

36from typing import Any 

37 

38try: 

39 from exec_runner import KubernetesExec, run_engine_command 

40except ImportError: # pragma: no cover — main.py qualified imports 

41 from langgraph.exec_runner import KubernetesExec, run_engine_command 

42 

43log = logging.getLogger(__name__) 

44 

45_RUNTIMES = ("bash", "python") 

46 

47 

48def _hard_failure(code: str, message: str, details: dict | None = None) -> Any: 

49 """Build the engine's RunHardFailure without a module-level circular import.""" 

50 try: 

51 from runner import RunHardFailure 

52 except ImportError: # pragma: no cover — main.py qualified imports 

53 from langgraph.runner import RunHardFailure 

54 return RunHardFailure(code, message, details) 

55 

56 

57def _runtime_command(runtime: str, path: Path) -> str: 

58 """The shell command that runs *path* with the declared runtime.""" 

59 interpreter = "bash" if runtime == "bash" else sys.executable 

60 return f"{shlex.quote(interpreter)} {shlex.quote(str(path))}" 

61 

62 

63def _in_cluster() -> bool: 

64 """Whether the engine runs inside a Kubernetes pod (in-cluster config).""" 

65 return bool(os.getenv("KUBERNETES_SERVICE_HOST")) 

66 

67 

68def _safe_step_token(step_id: str) -> str: 

69 """A filename-safe token that stays a direct child of `.stepfunction/`.""" 

70 token = re.sub(r"[^A-Za-z0-9_.-]", "_", step_id) 

71 if not token or token in (".", ".."): 

72 return "step" 

73 return token 

74 

75 

76def _write_params_file(root: Path, step_id: str, params: dict) -> Path: 

77 """Write the resolved params JSON to the shared-workspace params file.""" 

78 params_file = root / ".stepfunction" / f"{_safe_step_token(step_id)}.params.json" 

79 params_file.parent.mkdir(parents=True, exist_ok=True) 

80 params_file.write_text(json.dumps(params), encoding="utf-8") 

81 return params_file 

82 

83 

84def _remove_params_file(params_file: Path) -> None: 

85 """Best-effort cleanup of the params file — must never mask the run result.""" 

86 with contextlib.suppress(OSError): 

87 params_file.unlink(missing_ok=True) 

88 

89 

90def _run_in_sidecar(command: str, *, root: Path, step_id: str, params: dict) -> dict: 

91 """Run the entrypoint in the sidecar with params fed by a file redirect. 

92 

93 The K8s exec v4 client has no half-close, so streaming stdin can stall an 

94 entrypoint that reads to EOF; the engine and sidecar share the workspace 

95 volume, so the params JSON is written to a file the sidecar redirects in. 

96 """ 

97 params_file = _write_params_file(root, step_id, params) 

98 try: 

99 return KubernetesExec().run_step_function( 

100 f"{command} < {shlex.quote(str(params_file))}" 

101 ) 

102 finally: 

103 _remove_params_file(params_file) 

104 

105 

106def _parse_outputs(stdout: str, name: str) -> dict: 

107 """Parse the entrypoint's stdout as a JSON object (ADR-0012 — no silent default).""" 

108 try: 

109 outputs = json.loads(stdout) 

110 except (ValueError, TypeError) as exc: 

111 raise _hard_failure( 

112 "VALIDATION_ERROR", 

113 f"StepFunction '{name}': stdout is not valid JSON: {exc}", 

114 {"step_function": name, "stdout": stdout}, 

115 ) from exc 

116 if not isinstance(outputs, dict): 

117 raise _hard_failure( 

118 "VALIDATION_ERROR", 

119 f"StepFunction '{name}': stdout JSON must be an object, got " 

120 f"{type(outputs).__name__}", 

121 {"step_function": name, "stdout": stdout}, 

122 ) 

123 return outputs 

124 

125 

126def _declared_outputs(spec: dict) -> list[dict]: 

127 """The declared `spec.interface.outputs` entries (empty when undeclared).""" 

128 interface = spec.get("interface") or {} 

129 return interface.get("outputs") or [] 

130 

131 

132def _enforce_outputs( 

133 outputs: dict, declared: list[dict], *, name: str, step_id: str 

134) -> dict: 

135 """Select declared outputs from stdout, failing fast on contract violations.""" 

136 declared_names = {entry["name"] for entry in declared} 

137 for entry in declared: 

138 output_name = entry["name"] 

139 if not entry.get("optional") and output_name not in outputs: 

140 raise _hard_failure( 

141 "UNRESOLVED_OUTPUT", 

142 f"StepFunction '{name}' (step '{step_id}') is missing required " 

143 f"output '{output_name}'", 

144 {"step_id": step_id, "name": output_name}, 

145 ) 

146 unknown = [key for key in outputs if key not in declared_names] 

147 if unknown: 

148 raise _hard_failure( 

149 "VALIDATION_ERROR", 

150 f"StepFunction '{name}' (step '{step_id}') produced undeclared " 

151 f"output key(s): {', '.join(unknown)}", 

152 {"step_id": step_id, "keys": unknown}, 

153 ) 

154 return { 

155 entry["name"]: outputs[entry["name"]] 

156 for entry in declared 

157 if entry["name"] in outputs 

158 } 

159 

160 

161class StepFunctionExecutor: 

162 """Executes `kind: StepFunction` components over the manifest registry. 

163 

164 `root` is the consumer root the `entrypoint` paths are relative to; it 

165 defaults to the engine's PROJECT_ROOT (workspace root, ADR-0024). 

166 """ 

167 

168 def __init__(self, registry: Any, *, root: Path | None = None) -> None: 

169 self._registry = registry 

170 self._root = root 

171 

172 def has(self, name: str) -> bool: 

173 return self._registry.get("StepFunction", name) is not None 

174 

175 def instruction_for(self, name: str) -> str | None: 

176 """The registered StepFunction's `spec.agent_instructions`, or None. 

177 

178 Consumer StepFunctions own their per-step prompt text; a registered 

179 StepFunction without the field yields None so the build-time whitelist 

180 validator can fail fast (AGENT_INSTRUCTION_MISSING). 

181 """ 

182 manifest = self._registry.get("StepFunction", name) 

183 if manifest is None: 

184 return None 

185 return (manifest.get("spec") or {}).get("agent_instructions") 

186 

187 def params_for(self, name: str) -> list[dict] | None: 

188 """The registered StepFunction's declared `spec.interface.params`, else None. 

189 

190 Public accessor for the declared param interface so callers (usage 

191 instructions, the requested-calls validator) do not reach into the 

192 private registry. 

193 """ 

194 manifest = self._registry.get("StepFunction", name) 

195 if manifest is None: 

196 return None 

197 interface = (manifest.get("spec") or {}).get("interface") or {} 

198 params = interface.get("params") 

199 return params if isinstance(params, list) else None 

200 

201 def execute(self, name: str, step: dict, state: dict) -> dict: 

202 manifest = self._registry.get("StepFunction", name) 

203 if manifest is None: 

204 raise _hard_failure( 

205 "VALIDATION_ERROR", 

206 f"Unknown StepFunction '{name}'", 

207 {"step_function": name}, 

208 ) 

209 

210 spec = manifest.get("spec") or {} 

211 runtime = spec.get("runtime") 

212 entrypoint = spec.get("entrypoint") 

213 step_id = step.get("id", "?") 

214 if runtime not in _RUNTIMES: 

215 raise _hard_failure( 

216 "VALIDATION_ERROR", 

217 f"StepFunction '{name}': unsupported runtime {runtime!r} " 

218 f"(expected one of {_RUNTIMES})", 

219 {"step_function": name, "runtime": runtime}, 

220 ) 

221 if not isinstance(entrypoint, str) or not entrypoint.strip(): 

222 raise _hard_failure( 

223 "VALIDATION_ERROR", 

224 f"StepFunction '{name}': 'entrypoint' is required (ADR-0012)", 

225 {"step_function": name, "entrypoint": entrypoint}, 

226 ) 

227 

228 declared_outputs = _declared_outputs(spec) 

229 root = self._root or Path(os.getenv("PROJECT_ROOT", "/project")) 

230 params = step.get("params") or {} 

231 command = _runtime_command(runtime, root / entrypoint) 

232 if _in_cluster(): 

233 result = _run_in_sidecar(command, root=root, step_id=step_id, params=params) 

234 else: 

235 result = run_engine_command(command, cwd=str(root), stdin=json.dumps(params)) 

236 if result["exit_code"] != 0: 

237 raise _hard_failure( 

238 "EXECUTION_ERROR", 

239 f"StepFunction '{name}' (step '{step_id}') exited with " 

240 f"{result['exit_code']}: {result['stderr']}", 

241 { 

242 "step_function": name, 

243 "step_id": step_id, 

244 "exit_code": result["exit_code"], 

245 "stderr": result["stderr"], 

246 }, 

247 ) 

248 parsed = _parse_outputs(result["stdout"], name) 

249 return { 

250 **state, 

251 step_id: _enforce_outputs( 

252 parsed, declared_outputs, name=name, step_id=step_id 

253 ), 

254 }