@cosmicdrift/kumiko-framework 0.215.4 → 0.215.6

package/src/engine/feature-ast/extractors/ai-steps.ts

@@ -0,0 +1,366 @@
+import type { CallExpression, Node, ObjectLiteralExpression, SourceFile } from "ts-morph";
+import { SyntaxKind } from "ts-morph";
+import type {
+  AiClassifyPattern,
+  AiExtractPattern,
+  AiGeneratePattern,
+  AiStepOpaqueArgs,
+  AiStepPolicy,
+} from "../patterns";
+import type { SourceLocation } from "../source-location";
+import { sourceLocationFromNode } from "../source-location";
+import {
+  type ExtractOutput,
+  fail,
+  findFunctionLiteral,
+  isPlainObject,
+  isRawRefSentinel,
+  ok,
+  readDataLiteralNode,
+  readStringOrRaw,
+} from "./shared";
+
+type AiStepKind = AiGeneratePattern["kind"] | AiExtractPattern["kind"] | AiClassifyPattern["kind"];
+
+type AiStepCommonExtracted = {
+  readonly source: SourceLocation;
+  readonly argsSource?: AiStepOpaqueArgs;
+  readonly stepKey?: string | AiStepOpaqueArgs;
+  readonly promptKey?: string | AiStepOpaqueArgs;
+  readonly promptFallback?: string | AiStepOpaqueArgs;
+  readonly defaults?: AiStepPolicy | AiStepOpaqueArgs;
+  readonly paramsSchemaSource?: SourceLocation;
+};
+
+function resolveObjectLiteralArg(node: Node): ObjectLiteralExpression | undefined {
+  const direct = node.asKind(SyntaxKind.ObjectLiteralExpression);
+  if (direct) return direct;
+  const identifier = node.asKind(SyntaxKind.Identifier);
+  if (!identifier) return undefined;
+  const varDecl = node.getSourceFile().getVariableDeclaration(identifier.getText());
+  return varDecl?.getInitializer()?.asKind(SyntaxKind.ObjectLiteralExpression);
+}
+
+function readPropertyInitializer(
+  obj: ObjectLiteralExpression,
+  propertyName: string,
+): import("ts-morph").Expression | undefined {
+  const prop = obj.getProperty(propertyName);
+  if (!prop) return undefined;
+  const assign = prop.asKind(SyntaxKind.PropertyAssignment);
+  if (assign) return assign.getInitializer();
+  const shorthand = prop.asKind(SyntaxKind.ShorthandPropertyAssignment);
+  if (shorthand) return shorthand.getNameNode();
+  return undefined;
+}
+
+function readEditableStringProp(
+  obj: ObjectLiteralExpression,
+  propertyName: string,
+): string | AiStepOpaqueArgs | undefined {
+  const init = readPropertyInitializer(obj, propertyName);
+  if (!init) return undefined;
+  const value = readStringOrRaw(init);
+  if (value === undefined) return undefined;
+  return value;
+}
+
+function readEditableDefaults(
+  node: import("ts-morph").Expression,
+): AiStepPolicy | AiStepOpaqueArgs | undefined {
+  const value = readDataLiteralNode(node);
+  if (isRawRefSentinel(value)) return value;
+  if (!isPlainObject(value)) return undefined;
+  if (typeof value["enabled"] !== "boolean") return undefined;
+  if (!isPlainObject(value["params"])) return undefined;
+  return {
+    enabled: value["enabled"],
+    params: value["params"] as Record<string, unknown>,
+    ...(typeof value["providerId"] === "string" && { providerId: value["providerId"] }),
+    ...(typeof value["model"] === "string" && { model: value["model"] }),
+  };
+}
+
+function readClassifyActions(
+  node: import("ts-morph").Expression,
+): readonly { readonly type: string; readonly description: string }[] | undefined {
+  const value = readDataLiteralNode(node);
+  if (!Array.isArray(value)) return undefined;
+  const out: { type: string; description: string }[] = [];
+  for (const entry of value) {
+    if (!isPlainObject(entry)) return undefined;
+    if (typeof entry["type"] !== "string" || typeof entry["description"] !== "string") {
+      return undefined;
+    }
+    out.push({ type: entry["type"], description: entry["description"] });
+  }
+  return out;
+}
+
+function extractAiStepCommon(
+  call: CallExpression,
+  sourceFile: SourceFile,
+  kind: AiStepKind,
+): ExtractOutput<AiStepCommonExtracted> {
+  const source = sourceLocationFromNode(call, sourceFile);
+  const arg = call.getArguments()[0];
+  if (!arg) {
+    return fail(kind, source, "expected one argument object");
+  }
+
+  const obj = resolveObjectLiteralArg(arg);
+  if (!obj) {
+    if (isRawRefSentinel(readDataLiteralNode(arg))) {
+      return ok({
+        source,
+        argsSource: { __raw: arg.getText() },
+      });
+    }
+    return fail(
+      kind,
+      source,
+      "argument must be an inline object literal or a same-file const resolving to one",
+    );
+  }
+
+  const stepKey = readEditableStringProp(obj, "stepKey");
+  if (stepKey === undefined) {
+    return fail(kind, source, "missing `stepKey` property");
+  }
+  const promptKey = readEditableStringProp(obj, "promptKey");
+  if (promptKey === undefined) {
+    return fail(kind, source, "missing `promptKey` property");
+  }
+  const promptFallbackInit = readPropertyInitializer(obj, "promptFallback");
+  if (!promptFallbackInit) {
+    return fail(kind, source, "missing `promptFallback` property");
+  }
+  const promptFallback = readStringOrRaw(promptFallbackInit);
+  if (promptFallback === undefined) {
+    return fail(kind, source, "`promptFallback` must be a string literal or identifier reference");
+  }
+
+  const defaultsInit = readPropertyInitializer(obj, "defaults");
+  if (!defaultsInit) {
+    return fail(kind, source, "missing `defaults` property");
+  }
+  const defaults = readEditableDefaults(defaultsInit);
+  if (defaults === undefined) {
+    return fail(kind, source, "`defaults` could not be read as a StepPolicy literal");
+  }
+
+  const paramsSchemaInit = readPropertyInitializer(obj, "paramsSchema");
+  if (!paramsSchemaInit) {
+    return fail(kind, source, "missing `paramsSchema` property");
+  }
+
+  return ok({
+    source,
+    stepKey,
+    promptKey,
+    promptFallback,
+    defaults,
+    paramsSchemaSource: sourceLocationFromNode(paramsSchemaInit, sourceFile),
+  });
+}
+
+function requireResolvedCommon(
+  kind: AiStepKind,
+  pattern: AiStepCommonExtracted,
+): ExtractOutput<
+  Required<
+    Pick<
+      AiStepCommonExtracted,
+      "stepKey" | "promptKey" | "promptFallback" | "defaults" | "paramsSchemaSource"
+    >
+  > & { readonly source: SourceLocation }
+> {
+  if (pattern.argsSource) {
+    return fail(kind, pattern.source, "expected resolved inline object, got opaque args reference");
+  }
+  const { stepKey, promptKey, promptFallback, defaults, paramsSchemaSource, source } = pattern;
+  if (
+    stepKey === undefined ||
+    promptKey === undefined ||
+    promptFallback === undefined ||
+    defaults === undefined ||
+    paramsSchemaSource === undefined
+  ) {
+    return fail(kind, source, "resolved AI step is missing required header fields");
+  }
+  return ok({ source, stepKey, promptKey, promptFallback, defaults, paramsSchemaSource });
+}
+
+export function extractAiGenerate(
+  call: CallExpression,
+  sourceFile: SourceFile,
+): ExtractOutput<AiGeneratePattern> {
+  const common = extractAiStepCommon(call, sourceFile, "ai.generate");
+  if (common.kind === "error") return common;
+  if (common.pattern.argsSource) {
+    return ok({
+      kind: "ai.generate",
+      source: common.pattern.source,
+      argsSource: common.pattern.argsSource,
+    });
+  }
+
+  const arg = call.getArguments()[0];
+  const obj = arg ? resolveObjectLiteralArg(arg) : undefined;
+  if (!obj) {
+    return fail("ai.generate", common.pattern.source, "expected resolvable argument object");
+  }
+  const inputInit = readPropertyInitializer(obj, "input");
+  if (!inputInit) {
+    return fail("ai.generate", common.pattern.source, "missing `input` property");
+  }
+  const fn = findFunctionLiteral(inputInit);
+  if (!fn) {
+    return fail(
+      "ai.generate",
+      common.pattern.source,
+      "`input` must be an inline arrow function or function expression",
+    );
+  }
+
+  const resolved = requireResolvedCommon("ai.generate", common.pattern);
+  if (resolved.kind === "error") return resolved;
+  return ok({
+    kind: "ai.generate",
+    source: resolved.pattern.source,
+    stepKey: resolved.pattern.stepKey,
+    promptKey: resolved.pattern.promptKey,
+    promptFallback: resolved.pattern.promptFallback,
+    defaults: resolved.pattern.defaults,
+    paramsSchemaSource: resolved.pattern.paramsSchemaSource,
+    inputBody: sourceLocationFromNode(fn, sourceFile),
+  });
+}
+
+export function extractAiExtract(
+  call: CallExpression,
+  sourceFile: SourceFile,
+): ExtractOutput<AiExtractPattern> {
+  const common = extractAiStepCommon(call, sourceFile, "ai.extract");
+  if (common.kind === "error") return common;
+  if (common.pattern.argsSource) {
+    return ok({
+      kind: "ai.extract",
+      source: common.pattern.source,
+      argsSource: common.pattern.argsSource,
+    });
+  }
+
+  const arg = call.getArguments()[0];
+  const obj = arg ? resolveObjectLiteralArg(arg) : undefined;
+  if (!obj) {
+    return fail("ai.extract", common.pattern.source, "expected resolvable argument object");
+  }
+
+  const outputSchemaInit = readPropertyInitializer(obj, "outputSchema");
+  if (!outputSchemaInit) {
+    return fail("ai.extract", common.pattern.source, "missing `outputSchema` property");
+  }
+  const instructionsInit = readPropertyInitializer(obj, "instructions");
+  if (!instructionsInit) {
+    return fail("ai.extract", common.pattern.source, "missing `instructions` property");
+  }
+  const instructionsFn = findFunctionLiteral(instructionsInit);
+  if (!instructionsFn) {
+    return fail(
+      "ai.extract",
+      common.pattern.source,
+      "`instructions` must be an inline arrow function or function expression",
+    );
+  }
+
+  const documentInit = readPropertyInitializer(obj, "document");
+  let documentBody: SourceLocation | undefined;
+  if (documentInit) {
+    const documentFn = findFunctionLiteral(documentInit);
+    if (!documentFn) {
+      return fail(
+        "ai.extract",
+        common.pattern.source,
+        "`document` must be an inline arrow function or function expression",
+      );
+    }
+    documentBody = sourceLocationFromNode(documentFn, sourceFile);
+  }
+
+  const resolved = requireResolvedCommon("ai.extract", common.pattern);
+  if (resolved.kind === "error") return resolved;
+  return ok({
+    kind: "ai.extract",
+    source: resolved.pattern.source,
+    stepKey: resolved.pattern.stepKey,
+    promptKey: resolved.pattern.promptKey,
+    promptFallback: resolved.pattern.promptFallback,
+    defaults: resolved.pattern.defaults,
+    paramsSchemaSource: resolved.pattern.paramsSchemaSource,
+    outputSchemaSource: sourceLocationFromNode(outputSchemaInit, sourceFile),
+    instructionsBody: sourceLocationFromNode(instructionsFn, sourceFile),
+    ...(documentBody !== undefined && { documentBody }),
+  });
+}
+
+export function extractAiClassify(
+  call: CallExpression,
+  sourceFile: SourceFile,
+): ExtractOutput<AiClassifyPattern> {
+  const common = extractAiStepCommon(call, sourceFile, "ai.classify");
+  if (common.kind === "error") return common;
+  if (common.pattern.argsSource) {
+    return ok({
+      kind: "ai.classify",
+      source: common.pattern.source,
+      argsSource: common.pattern.argsSource,
+    });
+  }
+
+  const arg = call.getArguments()[0];
+  const obj = arg ? resolveObjectLiteralArg(arg) : undefined;
+  if (!obj) {
+    return fail("ai.classify", common.pattern.source, "expected resolvable argument object");
+  }
+
+  const actionsInit = readPropertyInitializer(obj, "actions");
+  if (!actionsInit) {
+    return fail("ai.classify", common.pattern.source, "missing `actions` property");
+  }
+  const actions = readClassifyActions(actionsInit);
+  if (actions === undefined) {
+    return fail(
+      "ai.classify",
+      common.pattern.source,
+      "`actions` must be an inline array of { type, description } objects",
+    );
+  }
+
+  const inputInit = readPropertyInitializer(obj, "input");
+  if (!inputInit) {
+    return fail("ai.classify", common.pattern.source, "missing `input` property");
+  }
+  const inputFn = findFunctionLiteral(inputInit);
+  if (!inputFn) {
+    return fail(
+      "ai.classify",
+      common.pattern.source,
+      "`input` must be an inline arrow function or function expression",
+    );
+  }
+
+  const resolved = requireResolvedCommon("ai.classify", common.pattern);
+  if (resolved.kind === "error") return resolved;
+  return ok({
+    kind: "ai.classify",
+    source: resolved.pattern.source,
+    stepKey: resolved.pattern.stepKey,
+    promptKey: resolved.pattern.promptKey,
+    promptFallback: resolved.pattern.promptFallback,
+    defaults: resolved.pattern.defaults,
+    paramsSchemaSource: resolved.pattern.paramsSchemaSource,
+    actions,
+    inputBody: sourceLocationFromNode(inputFn, sourceFile),
+  });
+}

package/src/engine/feature-ast/extractors/index.ts

@@ -1,3 +1,8 @@
+export {
+  extractAiClassify,
+  extractAiExtract,
+  extractAiGenerate,
+} from "./ai-steps";
 export {
   extractDefineEvent,
   extractNotification,

package/src/engine/feature-ast/parse.ts

@@ -24,7 +24,9 @@
 import type {
   ArrowFunction,
   CallExpression,
+  Expression,
   Node,
+  ObjectLiteralExpression,
   ParameterDeclaration,
   SourceFile,
 } from "ts-morph";
@@ -32,6 +34,9 @@ import { Project, SyntaxKind } from "ts-morph";
 
 import {
   type ExtractOutput,
+  extractAiClassify,
+  extractAiExtract,
+  extractAiGenerate,
   extractAuthClaims,
   extractClaimKey,
   extractConfig,
@@ -152,6 +157,8 @@ export function parseSourceFile(sourceFile: SourceFile): ParseResult {
   const errors: ParseError[] = [];
 
   walkSetupCallback(setupCallback.getBody(), registrarParamName, sourceFile, patterns, errors);
+  walkAiStepCalls(setupCallback.getBody(), registrarParamName, sourceFile, patterns, errors);
+  patterns.sort((a, b) => a.source.start.line - b.source.start.line);
 
   return { featureName, patterns, errors };
 }
@@ -380,6 +387,134 @@ function extractRegistrarMethodName(
   return propAccess.getName();
 }
 
+function readObjectPropertyInitializer(
+  obj: ObjectLiteralExpression,
+  propertyName: string,
+): Expression | undefined {
+  const prop = obj.getProperty(propertyName);
+  if (!prop) return undefined;
+  const assign = prop.asKind(SyntaxKind.PropertyAssignment);
+  if (assign) return assign.getInitializer();
+  const shorthand = prop.asKind(SyntaxKind.ShorthandPropertyAssignment);
+  if (shorthand) return shorthand.getNameNode();
+  return undefined;
+}
+
+function resolveSameFileObjectLiteralArg(node: Node): ObjectLiteralExpression | undefined {
+  const direct = node.asKind(SyntaxKind.ObjectLiteralExpression);
+  if (direct) return direct;
+  const identifier = node.asKind(SyntaxKind.Identifier);
+  if (!identifier) return undefined;
+  const varDecl = node.getSourceFile().getVariableDeclaration(identifier.getText());
+  return varDecl?.getInitializer()?.asKind(SyntaxKind.ObjectLiteralExpression);
+}
+
+function resolveStepsArrayRoot(stepsInit: Expression): Node | undefined {
+  const directArray = stepsInit.asKind(SyntaxKind.ArrayLiteralExpression);
+  if (directArray) return directArray;
+
+  const pipelineCall = stepsInit.asKind(SyntaxKind.CallExpression);
+  if (pipelineCall?.getExpression().getText() !== "stepsPipeline") {
+    return undefined;
+  }
+
+  const closureArg = pipelineCall.getArguments()[0];
+  if (!closureArg) return undefined;
+  const fn = findFunctionLiteral(closureArg);
+  if (!fn) return undefined;
+
+  const fnBody =
+    fn.asKind(SyntaxKind.ArrowFunction)?.getBody() ??
+    fn.asKind(SyntaxKind.FunctionExpression)?.getBody();
+  const exprBody = fnBody?.asKind(SyntaxKind.ArrayLiteralExpression);
+  if (exprBody) return exprBody;
+
+  if (fnBody?.isKind(SyntaxKind.Block)) {
+    for (const stmt of fnBody.getStatements()) {
+      const ret = stmt.asKind(SyntaxKind.ReturnStatement);
+      const retArray = ret?.getExpression()?.asKind(SyntaxKind.ArrayLiteralExpression);
+      if (retArray) return retArray;
+    }
+  }
+
+  return undefined;
+}
+
+function collectWorkflowStepArrayRoots(body: Node): Node[] {
+  const roots: Node[] = [];
+  for (const call of body.getDescendantsOfKind(SyntaxKind.CallExpression)) {
+    if (call.getExpression().getText() !== "defineWorkflow") continue;
+    const obj = resolveSameFileObjectLiteralArg(call.getArguments()[0] ?? call);
+    if (!obj) continue;
+    const stepsInit = readObjectPropertyInitializer(obj, "steps");
+    if (!stepsInit) continue;
+    const arrayRoot = resolveStepsArrayRoot(stepsInit);
+    if (arrayRoot) roots.push(arrayRoot);
+  }
+  return roots;
+}
+
+function walkAiStepCallsInNode(
+  node: Node,
+  sourceFile: SourceFile,
+  patterns: FeaturePattern[],
+  errors: ParseError[],
+): void {
+  for (const call of node.getDescendantsOfKind(SyntaxKind.CallExpression)) {
+    const callee = call.getExpression().getText();
+    switch (callee) {
+      case "aiGenerateStep": {
+        const result = extractAiGenerate(call, sourceFile);
+        if (result.kind === "pattern") patterns.push(result.pattern);
+        else errors.push(result.error);
+        break;
+      }
+      case "aiExtractStep": {
+        const result = extractAiExtract(call, sourceFile);
+        if (result.kind === "pattern") patterns.push(result.pattern);
+        else errors.push(result.error);
+        break;
+      }
+      case "aiClassifyStep": {
+        const result = extractAiClassify(call, sourceFile);
+        if (result.kind === "pattern") patterns.push(result.pattern);
+        else errors.push(result.error);
+        break;
+      }
+      default:
+        break;
+    }
+  }
+}
+
+function walkAiStepCalls(
+  body: Node,
+  registrarParamName: string,
+  sourceFile: SourceFile,
+  patterns: FeaturePattern[],
+  errors: ParseError[],
+  visitedWrapperDecls: Set<Node> = new Set(),
+): void {
+  for (const arrayRoot of collectWorkflowStepArrayRoots(body)) {
+    walkAiStepCallsInNode(arrayRoot, sourceFile, patterns, errors);
+  }
+
+  for (const call of body.getDescendantsOfKind(SyntaxKind.CallExpression)) {
+    const wrapper = resolveRegistrarWrapperCall(call, registrarParamName);
+    if (!wrapper) continue;
+    if (visitedWrapperDecls.has(wrapper.declNode)) continue;
+    visitedWrapperDecls.add(wrapper.declNode);
+    walkAiStepCalls(
+      wrapper.body,
+      wrapper.paramName,
+      wrapper.sourceFile,
+      patterns,
+      errors,
+      visitedWrapperDecls,
+    );
+  }
+}
+
 // =============================================================================
 // Internal — pattern extractors (skeleton, implementation in C1.5)
 // =============================================================================