r/src/messages.c

// retoor <retoor@molodetz.nl>

#include "messages.h"
#include "db.h"
#include <ctype.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#define MAX_CONTENT_LENGTH 1048570
#define MAX_TOOL_RESULT_LENGTH 104000

struct messages_t {
    struct json_object *array;
    char *session_id;
    db_handle db;
    bool loaded;
};

static bool is_valid_session_id(const char *session_id) {
    if (!session_id || !*session_id) return false;
    if (strlen(session_id) > 200) return false;

    for (const char *p = session_id; *p; p++) {
        if (!isalnum((unsigned char)*p) && *p != '-' && *p != '_' && *p != '.') {
            return false;
        }
    }
    return true;
}

static char *generate_default_session_id(void) {
    char *session_id = malloc(32);
    if (!session_id) return NULL;

    snprintf(session_id, 32, "session-%d", getppid());
    return session_id;
}

messages_handle messages_create(const char *session_id) {
    struct messages_t *msgs = calloc(1, sizeof(struct messages_t));
    if (!msgs) return NULL;

    msgs->array = json_object_new_array();
    if (!msgs->array) {
        free(msgs);
        return NULL;
    }

    if (session_id && is_valid_session_id(session_id)) {
        msgs->session_id = strdup(session_id);
    } else {
        msgs->session_id = generate_default_session_id();
    }

    if (!msgs->session_id) {
        json_object_put(msgs->array);
        free(msgs);
        return NULL;
    }

    msgs->db = db_open(NULL);

    return msgs;
}

void messages_destroy(messages_handle msgs) {
    if (!msgs) return;
    if (msgs->array) json_object_put(msgs->array);
    if (msgs->db) db_close(msgs->db);
    free(msgs->session_id);
    free(msgs);
}

r_status_t messages_set_session_id(messages_handle msgs, const char *session_id) {
    if (!msgs || !is_valid_session_id(session_id)) return R_ERROR_INVALID_ARG;

    free(msgs->session_id);
    msgs->session_id = strdup(session_id);
    msgs->loaded = false;
    return msgs->session_id ? R_SUCCESS : R_ERROR_OUT_OF_MEMORY;
}

const char *messages_get_session_id(messages_handle msgs) {
    return msgs ? msgs->session_id : NULL;
}

r_status_t messages_add(messages_handle msgs, const char *role, const char *content) {
    if (!msgs || !msgs->array || !role) return R_ERROR_INVALID_ARG;

    struct json_object *message = json_object_new_object();
    if (!message) return R_ERROR_OUT_OF_MEMORY;

    json_object_object_add(message, "role", json_object_new_string(role));

    if (content) {
        size_t len = strlen(content);
        if (len > MAX_CONTENT_LENGTH) len = MAX_CONTENT_LENGTH;
        json_object_object_add(message, "content",
            json_object_new_string_len(content, (int)len));
    }

    json_object_array_add(msgs->array, message);

    if (strcmp(role, "system") != 0) {
        messages_save(msgs);
    }

    return R_SUCCESS;
}

r_status_t messages_add_object(messages_handle msgs, struct json_object *message) {
    if (!msgs || !msgs->array || !message) return R_ERROR_INVALID_ARG;

    json_object_array_add(msgs->array, message);
    messages_save(msgs);
    return R_SUCCESS;
}

r_status_t messages_add_tool_call(messages_handle msgs, struct json_object *message) {
    if (!msgs || !msgs->array || !message) return R_ERROR_INVALID_ARG;

    json_object_array_add(msgs->array, message);
    messages_save(msgs);
    return R_SUCCESS;
}

r_status_t messages_add_tool_result(messages_handle msgs, const char *tool_call_id, const char *result) {
    if (!msgs || !msgs->array || !tool_call_id || !result) return R_ERROR_INVALID_ARG;

    struct json_object *message = json_object_new_object();
    if (!message) return R_ERROR_OUT_OF_MEMORY;

    json_object_object_add(message, "tool_call_id", json_object_new_string(tool_call_id));

    size_t len = strlen(result);
    if (len > MAX_TOOL_RESULT_LENGTH) len = MAX_TOOL_RESULT_LENGTH;
    json_object_object_add(message, "tool_result",
        json_object_new_string_len(result, (int)len));

    json_object_array_add(msgs->array, message);
    messages_save(msgs);
    return R_SUCCESS;
}

r_status_t messages_remove_last(messages_handle msgs) {
    if (!msgs || !msgs->array) return R_ERROR_INVALID_ARG;

    int size = json_object_array_length(msgs->array);
    if (size == 0) return R_ERROR_NOT_FOUND;

    json_object_array_del_idx(msgs->array, size - 1, 1);
    messages_save(msgs);
    return R_SUCCESS;
}

r_status_t messages_clear(messages_handle msgs) {
    if (!msgs) return R_ERROR_INVALID_ARG;

    if (msgs->array) json_object_put(msgs->array);
    msgs->array = json_object_new_array();
    if (!msgs->array) return R_ERROR_OUT_OF_MEMORY;

    messages_save(msgs);
    return R_SUCCESS;
}

r_status_t messages_save(messages_handle msgs) {
    if (!msgs || !msgs->array || !msgs->db) return R_ERROR_INVALID_ARG;

    char key[512];
    snprintf(key, sizeof(key), "session:%s", msgs->session_id);

    const char *json_str = json_object_to_json_string_ext(msgs->array, JSON_C_TO_STRING_PLAIN);
    if (!json_str) return R_ERROR_OUT_OF_MEMORY;

    return db_save_conversation(msgs->db, key, json_str);
}

r_status_t messages_load(messages_handle msgs) {
    if (!msgs || !msgs->db) return R_ERROR_INVALID_ARG;
    if (msgs->loaded) return R_SUCCESS;

    char key[512];
    snprintf(key, sizeof(key), "session:%s", msgs->session_id);

    char *data = NULL;
    r_status_t status = db_load_conversation(msgs->db, key, &data);
    if (status != R_SUCCESS || !data) {
        if (status == R_SUCCESS) msgs->loaded = true;
        return status == R_SUCCESS ? R_ERROR_NOT_FOUND : status;
    }

    struct json_object *loaded = json_tokener_parse(data);
    free(data);

    if (!loaded || !json_object_is_type(loaded, json_type_array)) {
        if (loaded) json_object_put(loaded);
        return R_ERROR_PARSE;
    }

    int len = json_object_array_length(loaded);
    for (int i = 0; i < len; i++) {
        struct json_object *msg = json_object_array_get_idx(loaded, i);
        struct json_object *role_obj;

        if (json_object_object_get_ex(msg, "role", &role_obj)) {
            const char *role = json_object_get_string(role_obj);
            if (role && strcmp(role, "system") != 0) {
                json_object_array_add(msgs->array, json_object_get(msg));
            }
        } else {
            json_object_array_add(msgs->array, json_object_get(msg));
        }
    }

    json_object_put(loaded);
    msgs->loaded = true;
    return R_SUCCESS;
}

struct json_object *messages_to_json(messages_handle msgs) {
    return msgs ? msgs->array : NULL;
}

char *messages_to_string(messages_handle msgs) {
    if (!msgs || !msgs->array) return NULL;
    return strdup(json_object_to_json_string_ext(msgs->array, JSON_C_TO_STRING_PRETTY));
}

int messages_count(messages_handle msgs) {
    if (!msgs || !msgs->array) return 0;
    return json_object_array_length(msgs->array);
}