tai/taicpp/tests/agent_memstore.h

#pragma once
// In-memory StorePort used by the agent tests. It implements the Store
// behaviour the covered scenarios exercise (sessions, systems, bots, secrets,
// redaction, schedules, audit, file records, events, turn state, records);
// FTS-backed searches are not emulated and throw.
#include "../src/agent/ports.h"
#include "../src/agent/text.h"
#include "../src/llm/helpers.h"
#include "../src/llm/py.h"
#include <algorithm>
#include <map>
#include <random>
namespace detz::agent::testing {

class MemStore : public StorePort {
 public:
  struct Event { std::string profile, role, kind, text; };
  struct Audit {
    int64_t id; std::string profile, ts, actor, action, path, message, tags;
    std::optional<int64_t> old_size, new_size;
    std::optional<std::string> old_text, new_text;
  };
  std::vector<Event> events;
  std::vector<Audit> audits;
  std::map<std::string, std::string> file_records;   // path -> content (kind file)
  std::map<std::string, Value> dir_defaults_log;

  std::string profile() const override { return profile_; }
  void set_profile(const std::string& name) override { profile_ = name; }
  bool seal_enabled() const override { return seal_; }
  bool fts_ok() const override { return false; }
  void set_seal_enabled(bool v) { seal_ = v; }

  std::string redact(const std::string& text, const std::string& profile) override {
    std::string out = text;
    auto& vault = secrets_[scope(profile)];
    std::vector<std::pair<std::string, std::string>> items;
    for (const auto& n : vault.order) {
      const auto& value = vault.values.at(n).value;
      if (py::len(value) >= 4) { items.emplace_back(n, value); }
    }
    std::stable_sort(items.begin(), items.end(), [](const auto& a, const auto& b) { return py::len(a.second) > py::len(b.second); });
    for (const auto& it : items) {
      if (out.find(it.second) != std::string::npos) { out = py::replace(out, it.second, "[redacted:" + it.first + "]"); }
    }
    return out;
  }
  void log_event(const std::string& profile, const std::string& role, const std::string& kind, const std::string& text) override {
    events.push_back({scope(profile), role, kind, text});
  }

  Value load_session(const std::string& name, const std::string& bot, const std::string& ctx) override {
    Value out = Value::array();
    auto it = sessions_.find(name + "|" + bot + "|" + ctx);
    if (it == sessions_.end()) { return out; }
    for (const auto& m : it->second.items()) {
      const Value& r = m.get("role");
      if (m.is_object() && r.is_string() && (r.str() == "user" || r.str() == "assistant" || r.str() == "tool")) { out.push(m); }
    }
    return out;
  }
  void save_session(const std::string& name, const Value& messages, const std::string& bot, const std::string& ctx) override {
    Value kept = Value::array();
    for (const auto& m : messages.items()) {
      if (!(m.get("role").is_string() && m.get("role").str() == "system")) { kept.push(m); }
    }
    sessions_[name + "|" + bot + "|" + ctx] = Value::loads(kept.dumps());
  }
  Value list_sessions(const std::string& profile, int64_t limit) override {
    Value out = Value::array();
    std::string want = scope(profile);
    std::vector<std::string> keys;
    for (const auto& kv : sessions_) { keys.push_back(kv.first); }
    std::sort(keys.begin(), keys.end());
    int64_t keep = std::max<int64_t>(1, std::min<int64_t>(limit, 50));
    int64_t shown = 0;
    for (const auto& key : keys) {
      size_t a = key.find('|');
      size_t b = key.find('|', a + 1);
      if (a == std::string::npos || b == std::string::npos || key.substr(0, a) != want) { continue; }
      if (shown++ >= keep) { break; }
      Value o = Value::object();
      o.set("bot", key.substr(a + 1, b - a - 1));
      o.set("sid", key.substr(b + 1));
      o.set("path", "(memstore)");
      o.set("mtime", 0.0);
      o.set("messages", static_cast<int>(sessions_.at(key).size()));
      o.set("goal", "");
      out.push(o);
    }
    return out;
  }
  std::optional<std::string> load_system(const std::string& name) override {
    auto it = systems_.find(name);
    if (it == systems_.end()) { return std::nullopt; }
    return it->second;
  }
  void save_system(const std::string& name, const std::string& text) override { systems_[name] = text; }
  std::optional<std::string> load_bot_system(const std::string& profile, const std::string& bot) override {
    if (bot == "main") { return load_system(profile); }
    auto it = bot_systems_.find(profile + "|" + bot);
    if (it == bot_systems_.end()) { return std::nullopt; }
    return it->second;
  }
  void save_bot_system(const std::string& profile, const std::string& bot, const std::string& text) override {
    bot_systems_[profile + "|" + bot] = text;
  }
  Value load_bots(const std::string& profile) override {
    auto it = bots_.find(profile);
    return it == bots_.end() ? Value::object() : Value::loads(it->second.dumps());
  }
  void save_bots(const std::string& profile, const Value& bots) override { bots_[profile] = Value::loads(bots.dumps()); }
  Value list_bots(const std::string& profile) override {
    Value bots = load_bots(profile);
    Value found = Value::array();
    Value main = Value::object();
    main.set("name", "main");
    main.set("nicknames", Value::array());
    found.push(main);
    std::vector<std::string> names;
    for (const auto& kv : bots.members()) { names.push_back(kv.first); }
    std::sort(names.begin(), names.end());
    for (const auto& n : names) {
      Value b = Value::object();
      b.set("name", n);
      const Value& nick = bots.get(n).get("nicknames");
      b.set("nicknames", nick.is_array() ? nick : Value::array());
      found.push(b);
    }
    return found;
  }
  std::optional<std::string> resolve_bot(const std::string& profile, const std::string& mention) override {
    std::string want = py::lower(py::strip(mention));
    if (want.empty() || want == "main") { return "main"; }
    Value bots = load_bots(profile);
    if (bots.has(want)) { return want; }
    std::vector<std::string> names;
    for (const auto& kv : bots.members()) { names.push_back(kv.first); }
    std::sort(names.begin(), names.end());
    for (const auto& n : names) {
      for (const auto& nick : bots.get(n).get("nicknames").items()) { if (nick.is_string() && nick.str() == want) { return n; } }
    }
    return std::nullopt;
  }
  Value load_turn_state(const std::string& profile, const std::string& bot, const std::string&) override {
    auto it = turns_.find(profile + "|" + bot);
    return it == turns_.end() ? Value::object() : it->second;
  }
  void save_turn_state(const std::string& profile, const std::string& bot, const Value& state, const std::string&) override {
    turns_[profile + "|" + bot] = Value::loads(state.dumps());
  }
  void save_dir_defaults(const std::string& directory, const std::optional<std::string>& provider,
                         const std::optional<std::string>& model, const std::optional<std::string>& bot,
                         const std::vector<std::string>& clear) override {
    Value o = Value::object();
    o.set("provider", provider ? Value(*provider) : Value());
    o.set("model", model ? Value(*model) : Value());
    o.set("bot", bot ? Value(*bot) : Value());
    Value c = Value::array();
    for (const auto& x : clear) { c.push(x); }
    o.set("clear", c);
    dir_defaults_log[directory] = o;
  }

  std::string add_record(const std::string& kind, const std::string& title, const std::string& content,
                         const std::vector<std::string>& tags, const std::string& profile) override {
    std::string id = "mem:" + py::uuid4_hex().substr(0, 16);
    Rec r{id, kind, title, content, {"record"}, 0};
    for (const auto& t : normalize_tags(tags)) { if (std::find(r.tags.begin(), r.tags.end(), t) == r.tags.end()) { r.tags.push_back(t); } }
    records_[scope(profile)].push_back(r);
    return id;
  }
  void add_chunks(const std::string& kind, const std::string& id, const std::vector<std::string>& texts, const std::string& profile) override {
    chunks_[scope(profile) + "|" + kind + "|" + id] = texts;
  }
  Value get_chunks_by_ids(const std::vector<int64_t>&, const std::string&) override { return Value::object(); }
  Value read_record(const std::string& id, int64_t offset, int64_t limit, const std::string& profile) override {
    for (auto& r : records_[scope(profile)]) {
      if (r.id != id) { continue; }
      int64_t total = static_cast<int64_t>(py::len(r.content));
      int64_t start = std::max<int64_t>(0, std::min(offset, total));
      int64_t end = std::min(total, start + std::max<int64_t>(1, limit));
      r.reads++;
      Value o = Value::object();
      o.set("id", id);
      o.set("slice", py::slice(r.content, start, end));
      o.set("start", start); o.set("end", end); o.set("total", total);
      Value t = Value::array();
      for (const auto& x : r.tags) { t.push(x); }
      o.set("tags", t);
      return o;
    }
    return Value();
  }
  Value search_records(const std::string&, const std::optional<std::string>&, const std::vector<std::string>&, int64_t, const std::string&) override {
    throw std::runtime_error("MemStore: search_records unsupported");
  }
  bool delete_record(const std::string& id, const std::string& profile) override {
    auto& v = records_[scope(profile)];
    for (auto it = v.begin(); it != v.end(); ++it) {
      if (it->id == id) { v.erase(it); return true; }
    }
    return false;
  }
  std::vector<std::string> item_tags(const std::string& item, const std::string& profile) override {
    for (const auto& r : records_[scope(profile)]) { if (r.id == item) { return r.tags; } }
    return {};
  }
  std::vector<std::pair<std::string, int64_t>> tag_counts(const std::string&, int64_t, const std::string&) override { return {}; }
  std::string add_edge(const std::string&, const std::string&, const Value&, const std::string&) override {
    throw std::runtime_error("MemStore: add_edge unsupported");
  }
  Value edges_for(const std::string&, const std::string&) override { return Value::array(); }
  bool has_node(const std::string&, const std::string&, const std::string&) override { return false; }
  std::string node_title(const std::string& node, const std::string&) override { return node; }
  Value traverse(const std::string&, int64_t, int64_t, const std::string&) override { return Value::array(); }
  Value fts_search(const std::string&, const std::vector<std::string>&, const std::string&, int64_t) override { return Value::array(); }
  Value events_like(const std::string&, const std::string&, int64_t) override { return Value::array(); }
  Value search_events(const std::string&, const std::string&, const std::vector<std::string>&) override { return Value::array(); }
  Value audit_search(const std::string&, const std::string&, int64_t) override { return Value::array(); }

  int64_t audit_event(const std::string& actor, const std::string& action, const std::string& path, const std::string& message,
                      const std::optional<std::string>& old_text, const std::optional<std::string>& new_text,
                      const std::vector<std::string>& tags, const std::optional<int64_t>& old_size,
                      const std::optional<int64_t>& new_size, const std::string& profile) override {
    Capped o = capped_text(old_text), n = capped_text(new_text);
    std::optional<std::string> ot = o.text, nt = n.text;
    std::optional<int64_t> os = o.size, ns = n.size;
    if (old_text.has_value() && old_size.has_value() && *old_size > static_cast<int64_t>(old_text->size())) {
      ot = py::head(*old_text, 20000) + trunc_marker(*old_size);
      os = old_size;
    }
    if (new_text.has_value() && new_size.has_value() && *new_size > static_cast<int64_t>(new_text->size())) {
      nt = py::head(*new_text, 20000) + trunc_marker(*new_size);
      ns = new_size;
    }
    Audit a;
    a.id = static_cast<int64_t>(audits.size()) + 1;
    a.profile = scope(profile);
    a.ts = py::utc_now_iso();
    a.actor = actor; a.action = action; a.path = path;
    a.message = py::head(message, 2000);
    a.tags = py::join(normalize_tags(tags), " ");
    a.old_size = os; a.new_size = ns; a.old_text = ot; a.new_text = nt;
    audits.push_back(a);
    return a.id;
  }
  Value audit_history(const std::optional<std::string>& path, const std::optional<std::string>& tag, int64_t limit,
                      const std::string& profile) override {
    Value out = Value::array();
    int64_t lim = std::max<int64_t>(1, std::min<int64_t>(limit, 100));
    for (auto it = audits.rbegin(); it != audits.rend() && static_cast<int64_t>(out.size()) < lim; ++it) {
      if (it->profile != scope(profile)) { continue; }
      if (path.has_value() && it->path != *path) { continue; }
      if (tag.has_value() && it->tags.find(*tag) == std::string::npos) { continue; }
      out.push(audit_row(*it, false));
    }
    return out;
  }
  Value audit_get(int64_t row_id, const std::string& profile) override {
    for (const auto& a : audits) { if (a.id == row_id && a.profile == scope(profile)) { return audit_row(a, true); } }
    return Value();
  }
  void upsert_file_record(const std::string& path, const std::string& content, const std::optional<int64_t>&, const std::string&) override {
    file_records[path] = content;
  }
  void delete_file_record(const std::string& path, const std::string&) override { file_records.erase(path); }

  std::optional<std::string> load_secret(const std::string& name, const std::string& profile) override {
    auto& v = secrets_[scope(profile)];
    auto it = v.values.find(name);
    if (it == v.values.end()) { return std::nullopt; }
    return it->second.value;
  }
  void save_secret(const std::string& name, const std::string& value, const Value& meta, const std::vector<std::string>& tags,
                   const std::string& profile) override {
    auto& v = secrets_[scope(profile)];
    if (!v.values.count(name)) { v.order.push_back(name); }
    std::vector<std::string> all = {"secret"};
    for (const auto& t : normalize_tags(tags)) { if (std::find(all.begin(), all.end(), t) == all.end()) { all.push_back(t); } }
    v.values[name] = {value, meta, all};
  }
  Value list_secret_infos(const std::string& profile) override {
    auto& v = secrets_[scope(profile)];
    std::vector<std::string> names = v.order;
    std::sort(names.begin(), names.end());
    Value out = Value::array();
    for (const auto& n : names) {
      Value o = Value::object();
      o.set("name", n);
      o.set("meta", v.values.at(n).meta);
      Value t = Value::array();
      std::vector<std::string> sorted_tags = v.values.at(n).tags;
      std::sort(sorted_tags.begin(), sorted_tags.end());
      for (const auto& x : sorted_tags) { t.push(x); }
      o.set("tags", t);
      o.set("updated", "");
      out.push(o);
    }
    return out;
  }
  void delete_secret(const std::string& name, const std::string& profile) override {
    auto& v = secrets_[scope(profile)];
    v.values.erase(name);
    v.order.erase(std::remove(v.order.begin(), v.order.end(), name), v.order.end());
  }

  int64_t add_schedule(const std::string& name, const std::string& prompt, const std::string& profile, int64_t every_sec,
                       const std::string& next_run, int64_t timeout, const std::vector<std::string>& tags) override {
    Value o = Value::object();
    int64_t id = ++sched_id_;
    o.set("id", id); o.set("name", name); o.set("prompt", prompt); o.set("profile", profile);
    o.set("every", every_sec); o.set("next_run", next_run); o.set("timeout", timeout);
    o.set("status", "pending"); o.set("last", ""); o.set("result", "");
    Value t = Value::array();
    t.push("schedule");
    for (const auto& x : tags) { t.push(x); }
    o.set("tags", t);
    schedules_.push_back(o);
    return id;
  }
  void remove_schedule(int64_t row_id, const std::string& profile) override {
    for (auto it = schedules_.begin(); it != schedules_.end(); ++it) {
      if (it->get("id").as_int() == row_id && it->get("profile").str() == scope(profile)) { schedules_.erase(it); return; }
    }
  }
  Value list_schedules(const std::string& profile) override {
    Value out = Value::array();
    std::vector<Value> rows;
    for (const auto& s : schedules_) { if (s.get("profile").str() == scope(profile)) { rows.push_back(s); } }
    std::stable_sort(rows.begin(), rows.end(), [](const Value& a, const Value& b) { return a.get("next_run").str() < b.get("next_run").str(); });
    for (const auto& r : rows) { out.push(r); }
    return out;
  }

  Value mcp_list_servers() override { return Value::array(); }
  Value mcp_get_server(const std::string&) override { return Value(); }
  void mcp_save_server(const Value&) override {}
  void mcp_remove_server(const std::string&) override {}
  void mcp_set_enabled(const std::string&, const std::string&, bool) override {}
  std::vector<std::string> mcp_enabled_profiles(const std::string&) override { return {}; }
  std::vector<std::string> mcp_server_profiles(const std::string&) override { return {}; }
  void mcp_store_tools(const std::string&, const Value&) override {}
  Value mcp_cached_tools(const std::string&) override { return Value::array(); }
  std::vector<std::string> mcp_tool_pins(const std::string&) override { return {}; }
  Value sql_rows(const std::string&, const Value&) override { throw std::runtime_error("MemStore: sql unsupported"); }

 private:
  struct Rec { std::string id, kind, title, content; std::vector<std::string> tags; int64_t reads; };
  struct Sec { std::string value; Value meta; std::vector<std::string> tags; };
  struct Vault { std::vector<std::string> order; std::map<std::string, Sec> values; };
  std::string scope(const std::string& p) const { return p.empty() ? profile_ : p; }
  Value audit_row(const Audit& a, bool full) const {
    Value o = Value::object();
    o.set("id", a.id); o.set("ts", a.ts); o.set("actor", a.actor); o.set("action", a.action); o.set("path", a.path);
    o.set("message", a.message);
    o.set("old_size", a.old_size ? Value(*a.old_size) : Value());
    o.set("new_size", a.new_size ? Value(*a.new_size) : Value());
    if (full) {
      o.set("old", a.old_text ? Value(*a.old_text) : Value());
      o.set("new", a.new_text ? Value(*a.new_text) : Value());
    }
    o.set("tags", a.tags);
    return o;
  }
  std::string profile_{"default"};
  bool seal_{true};
  std::map<std::string, Value> sessions_, turns_, bots_;
  std::map<std::string, std::string> systems_, bot_systems_;
  std::map<std::string, std::vector<Rec>> records_;
  std::map<std::string, std::vector<std::string>> chunks_;
  std::map<std::string, Vault> secrets_;
  std::vector<Value> schedules_;
  int64_t sched_id_{0};
};

}  // namespace detz::agent::testing