CST-agent/tests/test_registry.c

/* Tool registry: schema / prompt generation, profiles and /tools switches. */
#include "cJSON.h"
#include "cst_agent.h"
#include "cst_tools.h"

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

static int nfail;

#define EXPECT(cond, msg)                                                      \
  do {                                                                         \
    if (!(cond)) {                                                             \
      fprintf(stderr, "FAIL %s:%d: %s\n", __FILE__, __LINE__, msg);            \
      nfail++;                                                                 \
    }                                                                          \
  } while (0)

static int schema_has(cJSON *arr, const char *name) {
  const cJSON *t = NULL;
  cJSON_ArrayForEach(t, arr) {
    const cJSON *fn = cJSON_GetObjectItemCaseSensitive(t, "function");
    const cJSON *nm = cJSON_GetObjectItemCaseSensitive(fn, "name");
    if (cJSON_IsString(nm) && strcmp(nm->valuestring, name) == 0)
      return 1;
  }
  return 0;
}

static int echo_run(const cst_tool_ctx_t *ctx, const ctl_tool_call_t *c,
                    cst_tool_result_t *r, void *ud) {
  (void)ctx;
  const cst_tool_def_t *self = ud;
  char *j = cst_tool_args_json(self, c);
  r->out = j;
  r->out_len = j ? strlen(j) : 0;
  r->exit_code = 0;
  return 0;
}

static void dynamic_tools(void) {
  cst_agent_config_t cfg;
  memset(&cfg, 0, sizeof(cfg));
  int base = cst_tool_count();
  cst_tool_def_t d;
  memset(&d, 0, sizeof(d));
  d.name = "dyn_echo";
  d.desc = "Echo the arguments.";
  d.schema_json = "{\"type\":\"object\",\"properties\":{\"text\":{"
                  "\"type\":\"string\"},\"n\":{\"type\":\"integer\"},"
                  "\"opts\":{\"type\":\"object\"}},\"required\":[\"text\"]}";
  d.perm = CST_TP_READ;
  d.run_ex = echo_run;
  d.owner = "test:echo";
  EXPECT(cst_tool_register(&d) == 0, "register");
  EXPECT(cst_tool_register(&d) != 0, "duplicate name refused");
  EXPECT(cst_tool_count() == base + 1, "count grows");
  const cst_tool_def_t *t = cst_tool_find("dyn_echo");
  EXPECT(t && t != &d && strcmp(t->owner, "test:echo") == 0, "copied def");
  /* the runner gets its own def as ud for typing */
  ((cst_tool_def_t *)t)->ud = (void *)t;

  cJSON *sc = cst_tools_schema_cfg(&cfg);
  char *ss = cJSON_PrintUnformatted(sc);
  EXPECT(ss && strstr(ss, "\"name\":\"dyn_echo\"") &&
             strstr(ss, "\"required\":[\"text\"]"),
         "schema passes the JSON Schema through");
  free(ss);
  cJSON_Delete(sc);
  char *pr = cst_tools_prompt(&cfg);
  EXPECT(pr && strstr(pr, "- dyn_echo(text, n?, opts?): Echo the arguments."),
         "text prompt lists schema params");
  free(pr);

  ctl_tool_call_t c;
  memset(&c, 0, sizeof(c));
  snprintf(c.name, sizeof(c.name), "dyn_echo");
  cst_tool_ctx_t ctx;
  memset(&ctx, 0, sizeof(ctx));
  cst_tool_result_t r;
  memset(&r, 0, sizeof(r));
  cst_tool_run(&ctx, &c, &r);
  EXPECT(r.exit_code == 1 && r.out && strstr(r.out, "needs argument \"text\""),
         "schema required checked");
  cst_tool_result_free(&r);
  c.nargs = 3;
  snprintf(c.args[0].key, sizeof(c.args[0].key), "text");
  c.args[0].value = "42";
  c.args[0].value_len = 2;
  snprintf(c.args[1].key, sizeof(c.args[1].key), "n");
  c.args[1].value = "7";
  c.args[1].value_len = 1;
  snprintf(c.args[2].key, sizeof(c.args[2].key), "opts");
  c.args[2].value = "{\"a\":true}";
  c.args[2].value_len = strlen(c.args[2].value);
  memset(&r, 0, sizeof(r));
  cst_tool_run(&ctx, &c, &r);
  EXPECT(r.exit_code == 0 && r.out &&
             strcmp(r.out, "{\"text\":\"42\",\"n\":7,\"opts\":{\"a\":true}}") == 0,
         "args typed by schema");
  cst_tool_result_free(&r);

  cst_tool_set_disabled(&cfg, "dyn_echo", 1);
  cst_tool_set_disabled(&cfg, "grep", 1);
  EXPECT(!cst_tool_active(&cfg, t) &&
             !cst_tool_active(&cfg, cst_tool_find("grep")),
         "disable by name");
  cst_tool_set_disabled(&cfg, "dyn_echo", 0);
  EXPECT(cst_tool_active(&cfg, t) &&
             !cst_tool_active(&cfg, cst_tool_find("grep")),
         "re-enable one name");
  cst_tool_set_disabled(&cfg, "grep", 0);
  EXPECT(cfg.tools_off[0] == 0, "set empties");

  char nm[64];
  cst_tool_sanitize_name("mcp__my server__do.it", nm, sizeof(nm));
  EXPECT(strcmp(nm, "mcp__my_server__do_it") == 0, "sanitize");

  EXPECT(cst_tool_unregister_owner("test:echo") == 1 &&
             !cst_tool_find("dyn_echo") && cst_tool_count() == base,
         "unregister by owner");
}

int main(void) {
  dynamic_tools();
  cst_agent_config_t cfg;
  memset(&cfg, 0, sizeof(cfg));

  /* every tool has a runner and a well-formed parameter list */
  EXPECT(cst_tool_count() >= 13, "registry size");
  for (int i = 0; i < cst_tool_count(); i++) {
    const cst_tool_def_t *t = cst_tool_at(i);
    EXPECT(t->run != NULL, t->name);
    EXPECT((t->perm == CST_TP_WRITE) == (t->preview != NULL),
           "write tools (and only those) have previews");
    EXPECT(cst_tool_find(t->name) == t, "find by name");
  }

  cJSON *all = cst_tools_schema_cfg(&cfg);
  EXPECT(cJSON_GetArraySize(all) == cst_tool_count(), "default: every tool");
  EXPECT(schema_has(all, "edit_file") && schema_has(all, "grep") &&
             schema_has(all, "web_fetch"),
         "default schema has new tools");
  cJSON_Delete(all);

  cfg.tool_profile = CST_TOOLS_MINIMAL;
  cJSON *min = cst_tools_schema_cfg(&cfg);
  EXPECT(cJSON_GetArraySize(min) == 4, "minimal: four tools");
  EXPECT(schema_has(min, "read_file") && schema_has(min, "run_shell") &&
             schema_has(min, "write_file") && schema_has(min, "apply_patch"),
         "minimal: the original four");
  char *ms = cJSON_PrintUnformatted(min);
  EXPECT(ms && strstr(ms, "\"description\":\"Read a UTF-8 text file under the "
                          "workspace.\""),
         "minimal keeps the original wording");
  EXPECT(ms && !strstr(ms, "offset"), "minimal hides optional params");
  free(ms);
  cJSON_Delete(min);

  char *pmin = cst_agent_build_prompt(&cfg, CST_PROTO_NATIVE_EOS, NULL);
  EXPECT(pmin && strcmp(pmin, cst_agent_system_prompt_for(
                                  CST_PROTO_NATIVE_EOS)) == 0,
         "minimal prompt == legacy prompt");
  EXPECT(pmin && strstr(pmin, "ls/grep via run_shell"), "legacy workflow");
  free(pmin);
  char *pctl = cst_agent_build_prompt(&cfg, CST_PROTO_CTL, NULL);
  EXPECT(pctl && !strstr(pctl, "must be quoted"), "minimal ctl: no new rule");
  free(pctl);

  cfg.tool_profile = CST_TOOLS_DEFAULT;
  char *pdef = cst_agent_build_prompt(&cfg, CST_PROTO_CTL_HEREDOC,
                                      "Always run `make test`.");
  EXPECT(pdef && strstr(pdef, "- edit_file(path, old_string, new_string, "
                              "replace_all?):"),
         "default prompt lists edit_file with optional marker");
  EXPECT(pdef && strstr(pdef, "list_dir, glob, grep"), "default workflow");
  EXPECT(pdef && strstr(pdef, "must be quoted"), "ctl numeric rule");
  EXPECT(pdef && strstr(pdef, "Always run `make test`."), "project notes");
  free(pdef);

  /* /tools off: grep disappears everywhere */
  cst_tool_set_disabled(&cfg, "grep", 1);
  EXPECT(!cst_tool_active(&cfg, cst_tool_find("grep")), "grep off");
  cJSON *off = cst_tools_schema_cfg(&cfg);
  EXPECT(!schema_has(off, "grep") && schema_has(off, "glob"), "schema off");
  cJSON_Delete(off);
  char *names = cst_tools_names(&cfg);
  EXPECT(names && !strstr(names, "grep") && strstr(names, "glob"),
         "names off");
  free(names);
  char *poff = cst_agent_build_prompt(&cfg, CST_PROTO_NATIVE_EOS, NULL);
  EXPECT(poff && !strstr(poff, "- grep(") && strstr(poff, "list_dir, glob)"),
         "prompt off");
  free(poff);

  int prof = -1;
  EXPECT(cst_tool_profile_parse("minimal", &prof) == 0 &&
             prof == CST_TOOLS_MINIMAL,
         "parse minimal");
  EXPECT(cst_tool_profile_parse("bogus", &prof) != 0, "parse bogus");

  if (nfail) {
    fprintf(stderr, "test_registry: %d failure(s)\n", nfail);
    return 1;
  }
  printf("test_registry: OK\n");
  return 0;
}