250 lines
9.7 KiB
C
250 lines
9.7 KiB
C
#include "gem/server.h"
|
|
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <string.h>
|
|
#include <sys/socket.h>
|
|
#include <sys/wait.h>
|
|
#include <time.h>
|
|
#include <unistd.h>
|
|
|
|
#define CHECK(condition) do { if (!(condition)) { \
|
|
fprintf(stderr, "check failed at %s:%d: %s\n", __FILE__, __LINE__, #condition); return 1; \
|
|
} } while (0)
|
|
|
|
static int send_all_test(int fd, const char *data, size_t size)
|
|
{
|
|
while (size) {
|
|
ssize_t count = write(fd, data, size);
|
|
if (count <= 0) return -1;
|
|
data += count; size -= (size_t)count;
|
|
}
|
|
return 0;
|
|
}
|
|
|
|
static char *exchange(const struct gem_server_config *config, const char *request)
|
|
{
|
|
int sockets[2];
|
|
if (socketpair(AF_UNIX, SOCK_STREAM, 0, sockets) != 0) return NULL;
|
|
if (send_all_test(sockets[1], request, strlen(request)) != 0) {
|
|
close(sockets[0]); close(sockets[1]); return NULL;
|
|
}
|
|
(void)shutdown(sockets[1], SHUT_WR);
|
|
char error[256];
|
|
(void)gem_server_handle_connection(sockets[0], config, error, sizeof(error));
|
|
close(sockets[0]);
|
|
size_t size = 0, capacity = 4096;
|
|
char *response = malloc(capacity);
|
|
if (!response) { close(sockets[1]); return NULL; }
|
|
for (;;) {
|
|
if (size + 1024 + 1 > capacity) {
|
|
capacity *= 2;
|
|
void *memory = realloc(response, capacity);
|
|
if (!memory) { free(response); close(sockets[1]); return NULL; }
|
|
response = memory;
|
|
}
|
|
ssize_t count = read(sockets[1], response + size, capacity - size - 1);
|
|
if (count < 0) { free(response); close(sockets[1]); return NULL; }
|
|
if (!count) break;
|
|
size += (size_t)count;
|
|
}
|
|
close(sockets[1]); response[size] = '\0';
|
|
return response;
|
|
}
|
|
|
|
static int fake_generate(void *backend, const struct gem_chat_request *request,
|
|
gem_stream_emit_fn emit, void *user,
|
|
size_t *input_tokens, size_t *output_tokens,
|
|
const char **finish_reason, char *error, size_t error_size)
|
|
{
|
|
(void)backend; (void)error; (void)error_size;
|
|
if (!strstr(request->prompt, "<|turn>user\nhello<turn|>")) return -1;
|
|
struct gem_stream_event reasoning = {
|
|
.kind = GEM_STREAM_REASONING, .text = "brief thought", .text_size = 13,
|
|
};
|
|
struct gem_stream_event heartbeat = {.kind = GEM_STREAM_HEARTBEAT};
|
|
struct gem_stream_event text = {
|
|
.kind = GEM_STREAM_TEXT, .text = "hi\n", .text_size = 3,
|
|
};
|
|
struct gem_stream_event tool = {
|
|
.kind = GEM_STREAM_TOOL_CALL, .tool_id = "call_1", .tool_name = "shell",
|
|
.tool_arguments_json = "{\"cmd\":\"pwd\"}",
|
|
};
|
|
if (emit(&heartbeat, user) != 0 || emit(&reasoning, user) != 0 ||
|
|
emit(&text, user) != 0 || emit(&tool, user) != 0)
|
|
return -1;
|
|
*input_tokens = 11; *output_tokens = 7; *finish_reason = "tool_calls";
|
|
return 0;
|
|
}
|
|
|
|
static struct gem_server_config config(void)
|
|
{
|
|
return (struct gem_server_config){
|
|
.model_id = "gemma-4-31b", .context_length = 32768, .default_max_tokens = 128,
|
|
.sampler = {.temperature = 0.8f, .top_k = 40, .top_p = 0.95f,
|
|
.repeat_penalty = 1.0f, .repeat_last_n = 64, .seed = 1},
|
|
.generate = fake_generate,
|
|
};
|
|
}
|
|
|
|
static int endpoints(void)
|
|
{
|
|
struct gem_server_config server = config();
|
|
char *response = exchange(&server, "GET /health HTTP/1.1\r\nHost: localhost\r\n\r\n");
|
|
CHECK(response && strstr(response, "HTTP/1.1 200 OK") && strstr(response, "{\"status\":\"ok\"}"));
|
|
free(response);
|
|
response = exchange(&server, "GET /v1/models HTTP/1.1\r\nHost: localhost\r\n\r\n");
|
|
CHECK(response && strstr(response, "\"id\":\"gemma-4-31b\"") &&
|
|
strstr(response, "\"n_ctx\":32768"));
|
|
free(response);
|
|
response = exchange(&server, "GET /props?model=gemma HTTP/1.1\r\nHost: localhost\r\n\r\n");
|
|
CHECK(response && strstr(response, "\"vision\":false"));
|
|
free(response);
|
|
return 0;
|
|
}
|
|
|
|
static int stream(void)
|
|
{
|
|
struct gem_server_config server = config();
|
|
const char body[] =
|
|
"{\"model\":\"gemma-4-31b\",\"stream\":true,"
|
|
"\"stream_options\":{\"include_usage\":true},"
|
|
"\"messages\":[{\"role\":\"user\",\"content\":\"hello\"}]}";
|
|
char request[1024];
|
|
int size = snprintf(request, sizeof(request),
|
|
"POST /v1/chat/completions HTTP/1.1\r\nHost: localhost\r\n"
|
|
"Content-Type: application/json\r\nContent-Length: %zu\r\n\r\n%s",
|
|
sizeof(body) - 1, body);
|
|
CHECK(size > 0 && (size_t)size < sizeof(request));
|
|
char *response = exchange(&server, request);
|
|
CHECK(response && strstr(response, "Content-Type: text/event-stream"));
|
|
CHECK(strstr(response, "\"role\":\"assistant\""));
|
|
CHECK(strstr(response, ": keepalive\n\n"));
|
|
CHECK(strstr(response, "\"reasoning_content\":\"brief thought\""));
|
|
CHECK(strstr(response, "\"content\":\"hi\\n\""));
|
|
CHECK(strstr(response, "\"tool_calls\":[{\"index\":0,\"id\":\"call_1\""));
|
|
CHECK(strstr(response, "\"arguments\":\"{\\\"cmd\\\":\\\"pwd\\\"}\""));
|
|
CHECK(strstr(response, "\"finish_reason\":\"tool_calls\""));
|
|
CHECK(strstr(response, "\"prompt_tokens\":11,\"completion_tokens\":7,\"total_tokens\":18"));
|
|
CHECK(strstr(response, "data: [DONE]\n\n"));
|
|
free(response);
|
|
return 0;
|
|
}
|
|
|
|
static int completion(void)
|
|
{
|
|
struct gem_server_config server = config();
|
|
const char body[] =
|
|
"{\"model\":\"gemma-4-31b\",\"stream\":false,"
|
|
"\"messages\":[{\"role\":\"user\",\"content\":\"hello\"}]}";
|
|
char request[1024];
|
|
int size = snprintf(request, sizeof(request),
|
|
"POST /v1/chat/completions HTTP/1.1\r\nHost: localhost\r\n"
|
|
"Content-Type: application/json\r\nContent-Length: %zu\r\n\r\n%s",
|
|
sizeof(body) - 1, body);
|
|
CHECK(size > 0 && (size_t)size < sizeof(request));
|
|
char *response = exchange(&server, request);
|
|
CHECK(response && strstr(response, "HTTP/1.1 200 OK"));
|
|
CHECK(strstr(response, "Content-Type: application/json"));
|
|
CHECK(strstr(response, "\"object\":\"chat.completion\""));
|
|
CHECK(strstr(response, "\"role\":\"assistant\",\"content\":\"hi\\n\""));
|
|
CHECK(strstr(response, "\"reasoning_content\":\"brief thought\""));
|
|
CHECK(strstr(response, "\"tool_calls\":[{\"id\":\"call_1\""));
|
|
CHECK(strstr(response, "\"arguments\":\"{\\\"cmd\\\":\\\"pwd\\\"}\""));
|
|
CHECK(strstr(response, "\"finish_reason\":\"tool_calls\""));
|
|
CHECK(strstr(response, "\"prompt_tokens\":11,\"completion_tokens\":7,\"total_tokens\":18"));
|
|
CHECK(!strstr(response, "keepalive"));
|
|
CHECK(!strstr(response, "data: "));
|
|
free(response);
|
|
return 0;
|
|
}
|
|
|
|
static int cancellation_generate(void *backend, const struct gem_chat_request *request,
|
|
gem_stream_emit_fn emit, void *user,
|
|
size_t *input_tokens, size_t *output_tokens,
|
|
const char **finish_reason, char *error, size_t error_size)
|
|
{
|
|
(void)request; (void)input_tokens; (void)output_tokens; (void)finish_reason;
|
|
int *cancelled = backend;
|
|
struct gem_stream_event heartbeat = {.kind = GEM_STREAM_HEARTBEAT};
|
|
const struct timespec delay = {.tv_nsec = 1000000};
|
|
for (size_t attempt = 0; attempt < 5000; ++attempt) {
|
|
if (emit(&heartbeat, user) != 0) {
|
|
*cancelled = 1;
|
|
(void)snprintf(error, error_size, "client disconnected");
|
|
return -1;
|
|
}
|
|
(void)nanosleep(&delay, NULL);
|
|
}
|
|
return -1;
|
|
}
|
|
|
|
static int cancellation(void)
|
|
{
|
|
int sockets[2];
|
|
CHECK(socketpair(AF_UNIX, SOCK_STREAM, 0, sockets) == 0);
|
|
pid_t reader = fork();
|
|
CHECK(reader >= 0);
|
|
if (reader == 0) {
|
|
close(sockets[0]);
|
|
const char body[] =
|
|
"{\"model\":\"gemma-4-31b\",\"stream\":true,"
|
|
"\"messages\":[{\"role\":\"user\",\"content\":\"hello\"}]}";
|
|
char request[1024];
|
|
int size = snprintf(request, sizeof(request),
|
|
"POST /v1/chat/completions HTTP/1.1\r\nHost: localhost\r\n"
|
|
"Content-Type: application/json\r\nContent-Length: %zu\r\n\r\n%s",
|
|
sizeof(body) - 1, body);
|
|
if (size <= 0 || (size_t)size >= sizeof(request) ||
|
|
send_all_test(sockets[1], request, (size_t)size) != 0)
|
|
_exit(1);
|
|
char response[4096] = {0};
|
|
size_t used = 0;
|
|
while (used + 1 < sizeof(response)) {
|
|
ssize_t count = read(sockets[1], response + used,
|
|
sizeof(response) - used - 1);
|
|
if (count <= 0) _exit(1);
|
|
used += (size_t)count;
|
|
response[used] = '\0';
|
|
if (strstr(response, "\"role\":\"assistant\"")) break;
|
|
}
|
|
close(sockets[1]);
|
|
_exit(0);
|
|
}
|
|
close(sockets[1]);
|
|
int cancelled = 0;
|
|
struct gem_server_config server = config();
|
|
server.backend = &cancelled;
|
|
server.generate = cancellation_generate;
|
|
char error[256] = {0};
|
|
CHECK(gem_server_handle_connection(sockets[0], &server, error, sizeof(error)) != 0);
|
|
close(sockets[0]);
|
|
int reader_status = 0;
|
|
CHECK(waitpid(reader, &reader_status, 0) == reader);
|
|
CHECK(WIFEXITED(reader_status) && WEXITSTATUS(reader_status) == 0);
|
|
CHECK(cancelled == 1);
|
|
CHECK(strstr(error, "client disconnected"));
|
|
return 0;
|
|
}
|
|
|
|
static int errors(void)
|
|
{
|
|
struct gem_server_config server = config();
|
|
const char request[] = "POST /v1/chat/completions HTTP/1.1\r\nHost: x\r\nContent-Length: 2\r\n\r\n{}";
|
|
char *response = exchange(&server, request);
|
|
CHECK(response && strstr(response, "HTTP/1.1 400 Bad Request"));
|
|
CHECK(strstr(response, "invalid_request_error"));
|
|
free(response);
|
|
return 0;
|
|
}
|
|
|
|
int main(void)
|
|
{
|
|
CHECK(endpoints() == 0);
|
|
CHECK(stream() == 0);
|
|
CHECK(completion() == 0);
|
|
CHECK(cancellation() == 0);
|
|
CHECK(errors() == 0);
|
|
puts("server tests passed");
|
|
return 0;
|
|
}
|